diff --git a/api.go b/api.go index 7c11a4e3e..4483e0db9 100644 --- a/api.go +++ b/api.go @@ -1661,6 +1661,38 @@ func (api *API) TranslateIDs(ctx context.Context, r io.Reader) (_ []byte, err er return buf, nil } +// FindIndexKeys looks up column keys in the index, mapping them to IDs. +// If a key does not exist, it will be absent from the resulting map. +func (api *API) FindIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) { + return api.cluster.findIndexKeys(ctx, index, keys...) +} + +// FindFieldKeys looks up keys in a field, mapping them to IDs. +// If a key does not exist, it will be absent from the resulting map. +func (api *API) FindFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) { + f := api.holder.Field(index, field) + if f == nil { + return nil, errors.Wrapf(ErrFieldNotFound, "finding keys for field %q", field) + } + return api.cluster.findFieldKeys(ctx, f, keys...) +} + +// CreateIndexKeys looks up column keys in the index, mapping them to IDs. +// If a key does not exist, it will be created. +func (api *API) CreateIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) { + return api.cluster.createIndexKeys(ctx, index, keys...) +} + +// CreateFieldKeys looks up keys in a field, mapping them to IDs. +// If a key does not exist, it will be created. +func (api *API) CreateFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) { + f := api.holder.Field(index, field) + if f == nil { + return nil, errors.Wrapf(ErrFieldNotFound, "finding keys for field %q", field) + } + return api.cluster.createFieldKeys(ctx, f, keys...) +} + // PrimaryReplicaNodeURL returns the URL of the cluster's primary replica. func (api *API) PrimaryReplicaNodeURL() url.URL { node := api.cluster.PrimaryReplicaNode() diff --git a/boltdb/translate.go b/boltdb/translate.go index d37dae0d9..7e0807889 100644 --- a/boltdb/translate.go +++ b/boltdb/translate.go @@ -190,6 +190,79 @@ func (s *TranslateStore) TranslateKeys(keys []string, writable bool) ([]uint64, return s.translateKeys(keys, writable) } +// FindKeys looks up the ID for each key. +// Keys are not created if they do not exist. +// Missing keys are not considered errors, so the length of the result may be less than that of the input. +func (s *TranslateStore) FindKeys(keys ...string) (map[string]uint64, error) { + result := make(map[string]uint64, len(keys)) + err := s.db.View(func(tx *bolt.Tx) error { + bkt := tx.Bucket(bucketKeys) + if bkt == nil { + return errors.Errorf(errFmtTranslateBucketNotFound, bucketKeys) + } + for _, key := range keys { + id, _ := findIDByKey(bkt, key) + if id == 0 { + // The key does not exist. + continue + } + + result[key] = id + } + return nil + }) + if err != nil { + return nil, err + } + + return result, nil +} + +// CreateKeys maps all keys to IDs, creating the IDs if they do not exist. +// If the translator is read-only, this will return an error. +func (s *TranslateStore) CreateKeys(keys ...string) (map[string]uint64, error) { + if s.ReadOnly() { + return nil, pilosa.ErrTranslateStoreReadOnly + } + + written := false + result := make(map[string]uint64, len(keys)) + err := s.db.Update(func(tx *bolt.Tx) error { + keyBucket := tx.Bucket(bucketKeys) + if keyBucket == nil { + return errors.Errorf(errFmtTranslateBucketNotFound, bucketKeys) + } + idBucket := tx.Bucket(bucketIDs) + if keyBucket == nil { + return errors.Errorf(errFmtTranslateBucketNotFound, bucketIDs) + } + for _, key := range keys { + id, boltKey := findIDByKey(keyBucket, key) + if id == 0 { + // The key does not exist. + id = pilosa.GenerateNextPartitionedID(s.index, maxID(tx), s.partitionID, s.partitionN) + if err := keyBucket.Put(boltKey, u64tob(id)); err != nil { + return err + } else if err := idBucket.Put(u64tob(id), boltKey); err != nil { + return err + } + written = true + } + + result[key] = id + } + return nil + }) + if err != nil { + return nil, err + } + if written { + s.notifyWrite() + } + + return result, nil +} + func (s *TranslateStore) translateKeys(keys []string, writable bool) ([]uint64, error) { ids := make([]uint64, 0, len(keys)) diff --git a/boltdb/translate_test.go b/boltdb/translate_test.go index 73b4e5345..b6827e89c 100644 --- a/boltdb/translate_test.go +++ b/boltdb/translate_test.go @@ -92,6 +92,46 @@ func TestTranslateStore_TranslateKeys(t *testing.T) { } } +func TestTranslateStore_CreateKeys(t *testing.T) { + s := MustOpenNewTranslateStore() + defer MustCloseTranslateStore(s) + + ids, err := s.CreateKeys("abc", "abc") + if err != nil { + t.Fatal(err) + } else if _, ok := ids["abc"]; !ok { + t.Fatalf(`missing "abc"; got %v`, ids) + } else if len(ids) > 1 { + t.Fatalf("expected one key, got %d in %v", len(ids), ids) + } + + // Ensure different keys translate to different IDs. + ids1, err := s.CreateKeys("foo", "bar") + if err != nil { + t.Fatal(err) + } else if foo, bar := ids1["foo"], ids1["bar"]; foo == bar { + t.Fatalf(`"foo" and "bar" map back to the same ID %d`, foo) + } + + // Ensure retranslation returns original IDs. + if ids, err := s.CreateKeys("bar", "foo"); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(ids, ids1) { + t.Fatalf("retranslation produced result %v which is different from original translation %v", ids, ids1) + } + + // Ensure retranslating with existing and non-existing keys returns correctly. + if ids, err := s.CreateKeys("foo", "baz", "bar"); err != nil { + t.Fatal(err) + } else if got, want := ids["foo"], ids1["foo"]; got != want { + t.Fatalf(`mismatched ID %d for "foo" (previously %d)`, got, want) + } else if _, ok := ids["baz"]; !ok { + t.Fatalf(`missing translation for "baz"; got %v`, ids) + } else if got, want := ids["bar"], ids1["bar"]; got != want { + t.Fatalf(`mismatched ID %d for "bar" (previously %d)`, got, want) + } +} + func TestTranslateStore_ReadKey(t *testing.T) { s := MustOpenNewTranslateStore() defer MustCloseTranslateStore(s) @@ -216,6 +256,84 @@ func TestTranslateStore_TranslateIDs(t *testing.T) { } } +func TestTranslateStore_FindKeys(t *testing.T) { + cases := []struct { + name string + data []string + lookup []string + }{ + { + name: "All", + data: []string{"plugh", "xyzzy", "h"}, + lookup: []string{"plugh", "xyzzy", "h"}, + }, + { + name: "Extra", + data: []string{"plugh", "xyzzy", "h"}, + lookup: []string{"plugh", "xyzzy", "h", "65"}, + }, + { + name: "None", + data: []string{"a", "b", "c"}, + lookup: []string{"d", "e"}, + }, + { + name: "Empty", + lookup: []string{"h"}, + }, + { + name: "LookupNothing", + }, + } + + for _, c := range cases { + c := c + t.Run(c.name, func(t *testing.T) { + s := MustOpenNewTranslateStore() + defer MustCloseTranslateStore(s) + + var naiveMap map[string]uint64 + if c.data != nil { + // Load in key data. + keys := c.data + ids, err := s.TranslateKeys(keys, true) + if err != nil { + t.Errorf("failed to import keys: %v", err) + return + } + if len(ids) != len(keys) { + t.Errorf("mapped %d keys to %d ids", len(keys), len(ids)) + return + } + naiveMap = make(map[string]uint64, len(keys)) + for i, key := range keys { + naiveMap[key] = ids[i] + } + } + + // Compute expected lookup result. + result := map[string]uint64{} + for _, key := range c.lookup { + id, ok := naiveMap[key] + if !ok { + // The key is expected to be missing. + continue + } + + result[key] = id + } + + // Find the keys. + found, err := s.FindKeys(c.lookup...) + if err != nil { + t.Errorf("failed to find keys: %v", err) + } else if !reflect.DeepEqual(result, found) { + t.Errorf("expected %v but found %v", result, found) + } + }) + } +} + func TestTranslateStore_EntryReader(t *testing.T) { t.Run("OK", func(t *testing.T) { s := MustOpenNewTranslateStore() diff --git a/client.go b/client.go index 42ec51ba0..28e79679c 100644 --- a/client.go +++ b/client.go @@ -92,11 +92,17 @@ type InternalQueryClient interface { // Trasnlate keys on the particular node. The parameter writable informs TranslateStore if we can generate a new ID if any of keys does not exist. TranslateKeysNode(ctx context.Context, uri *URI, index, field string, keys []string, writable bool) ([]uint64, error) TranslateIDsNode(ctx context.Context, uri *URI, index, field string, id []uint64) ([]string, error) + + FindIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) + FindFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) + + CreateIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) + CreateFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) } type nopInternalQueryClient struct{} -func (n *nopInternalQueryClient) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) { +func (n nopInternalQueryClient) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) { return nil, nil } @@ -108,15 +114,31 @@ func (n nopInternalQueryClient) TranslateIDsNode(ctx context.Context, uri *URI, return nil, nil } -func newNopInternalQueryClient() *nopInternalQueryClient { - return &nopInternalQueryClient{} +func (n nopInternalQueryClient) FindIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n nopInternalQueryClient) FindFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n nopInternalQueryClient) CreateIndexKeysNode(ctx context.Context, uri *URI, index string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n nopInternalQueryClient) CreateFieldKeysNode(ctx context.Context, uri *URI, index string, field string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func newNopInternalQueryClient() nopInternalQueryClient { + return nopInternalQueryClient{} } var _ InternalQueryClient = newNopInternalQueryClient() //=============== -type nopInternalClient struct{} +type nopInternalClient struct{ nopInternalQueryClient } func newNopInternalClient() nopInternalClient { return nopInternalClient{} @@ -144,15 +166,6 @@ func (n nopInternalClient) Nodes(ctx context.Context) ([]*Node, error) { func (n nopInternalClient) Query(ctx context.Context, index string, queryRequest *QueryRequest) (*QueryResponse, error) { return nil, nil } -func (n nopInternalClient) QueryNode(ctx context.Context, uri *URI, index string, queryRequest *QueryRequest) (*QueryResponse, error) { - return nil, nil -} -func (n nopInternalClient) TranslateKeysNode(ctx context.Context, uri *URI, index, field string, keys []string, writable bool) ([]uint64, error) { - return nil, nil -} -func (n nopInternalClient) TranslateIDsNode(ctx context.Context, uri *URI, index, field string, ids []uint64) ([]string, error) { - return nil, nil -} func (n nopInternalClient) Import(ctx context.Context, index, field string, shard uint64, bits []Bit, opts ...ImportOption) error { return nil } diff --git a/cluster.go b/cluster.go index 3905c91d1..534ec8429 100644 --- a/cluster.go +++ b/cluster.go @@ -2324,17 +2324,6 @@ func (c *cluster) setStatic(hosts []string) error { return nil } -// translateFieldKey gets a single key from translateFieldKeys. -func (c *cluster) translateFieldKey(ctx context.Context, field *Field, key string, writable bool) (uint64, error) { - ids, err := c.translateFieldKeys(ctx, field, []string{key}, writable) - if err != nil { - return 0, err - } else if len(ids) == 0 { - return 0, nil - } - return ids[0], nil -} - // translateFieldKeys is basically a wrapper around // field.TranslateStore().TranslateKey(key), but in // the case where the local node is not coordinator, then this method will forward the translation @@ -2359,6 +2348,133 @@ func (c *cluster) translateFieldKeys(ctx context.Context, field *Field, keys []s return ids, nil } +func (c *cluster) findFieldKeys(ctx context.Context, field *Field, keys ...string) (map[string]uint64, error) { + if idx := field.ForeignIndex(); idx != "" { + // The field uses foreign index keys. + // Therefore, the field keys are actually column keys on a different index. + return c.findIndexKeys(ctx, idx, keys...) + } + + if !field.Keys() { + return nil, errors.Wrap(ErrTranslatingKeyNotFound, "field is not keyed") + } + + // Attempt to find the keys locally. + localTranslations, err := field.TranslateStore().FindKeys(keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) locally", field.Index(), field.Name(), keys) + } + + // Check for missing keys. + var missing []string + if len(keys) > len(localTranslations) { + // There are either duplicate keys or missing keys. + // This should work either way. + missing = make([]string, 0, len(keys)-len(localTranslations)) + for _, k := range keys { + _, found := localTranslations[k] + if !found { + missing = append(missing, k) + } + } + } else if len(localTranslations) > len(keys) { + panic(fmt.Sprintf("more translations than keys! translation count=%v, key count=%v", len(localTranslations), len(keys))) + } + if len(missing) == 0 { + // All keys were available locally. + return localTranslations, nil + } + + // It is possible that the missing keys exist, but have not been synced to the local replica. + coordinator := c.coordinatorNode() + if coordinator == nil { + return nil, errors.Errorf("translating field(%s/%s) keys(%v) - cannot find coordinator node", field.Index(), field.Name(), keys) + } + if c.Node.ID == coordinator.ID { + // The local copy is the authoritative copy. + return localTranslations, nil + } + + // Forward the missing keys to the coordinator. + // The coordinator has the authoritative copy. + remoteTranslations, err := c.InternalClient.FindFieldKeysNode(ctx, &coordinator.URI, field.Index(), field.Name(), missing...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) remotely", field.Index(), field.Name(), keys) + } + + // Merge the remote translations into the local translations. + translations := localTranslations + for key, id := range remoteTranslations { + translations[key] = id + } + + return translations, nil +} + +func (c *cluster) createFieldKeys(ctx context.Context, field *Field, keys ...string) (map[string]uint64, error) { + if idx := field.ForeignIndex(); idx != "" { + // The field uses foreign index keys. + // Therefore, the field keys are actually column keys on a different index. + return c.createIndexKeys(ctx, idx, keys...) + } + + if !field.Keys() { + return nil, errors.Wrap(ErrTranslatingKeyNotFound, "field is not keyed") + } + + // The coordinator is the only node that can create field keys, since it owns the authoritative copy. + coordinator := c.coordinatorNode() + if coordinator == nil { + return nil, errors.Errorf("translating field(%s/%s) keys(%v) - cannot find coordinator node", field.Index(), field.Name(), keys) + } + if c.Node.ID == coordinator.ID { + // The local copy is the authoritative copy. + return field.TranslateStore().CreateKeys(keys...) + } + + // Attempt to find the keys locally. + // They cannot be created locally, but skipping keys that exist can reduce network usage. + localTranslations, err := field.TranslateStore().FindKeys(keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) locally", field.Index(), field.Name(), keys) + } + + // Check for missing keys. + var missing []string + if len(keys) > len(localTranslations) { + // There are either duplicate keys or missing keys. + // This should work either way. + missing = make([]string, 0, len(keys)-len(localTranslations)) + for _, k := range keys { + _, found := localTranslations[k] + if !found { + missing = append(missing, k) + } + } + } else if len(localTranslations) > len(keys) { + panic(fmt.Sprintf("more translations than keys! translation count=%v, key count=%v", len(localTranslations), len(keys))) + } + if len(missing) == 0 { + // All keys exist locally. + // There is no need to create anything. + return localTranslations, nil + } + + // Forward the missing keys to the coordinator to be created. + remoteTranslations, err := c.InternalClient.CreateFieldKeysNode(ctx, &coordinator.URI, field.Index(), field.Name(), missing...) + if err != nil { + return nil, errors.Wrapf(err, "translating field(%s/%s) keys(%v) remotely", field.Index(), field.Name(), keys) + } + + // Merge the remote translations into the local translations. + translations := localTranslations + for key, id := range remoteTranslations { + translations[key] = id + } + + return translations, nil +} + func (c *cluster) translateIndexKey(ctx context.Context, indexName string, key string, writable bool) (uint64, error) { keyMap, err := c.translateIndexKeySet(ctx, indexName, map[string]struct{}{key: struct{}{}}, writable) if err != nil { @@ -2450,6 +2566,208 @@ func (c *cluster) translateIndexKeySet(ctx context.Context, indexName string, ke return keyMap, nil } +func (c *cluster) findIndexKeys(ctx context.Context, indexName string, keys ...string) (map[string]uint64, error) { + done := ctx.Done() + + idx := c.holder.Index(indexName) + if idx == nil { + return nil, ErrIndexNotFound + } + + // Split keys by partition. + keysByPartition := make(map[int][]string, c.partitionN) + for _, key := range keys { + partitionID := c.keyPartition(indexName, key) + keysByPartition[partitionID] = append(keysByPartition[partitionID], key) + } + + // TODO: use local replicas to short-circuit network traffic + + // Group keys by node. + keysByNode := make(map[*Node][]string) + for partitionID, keys := range keysByPartition { + // Find the primary node for this partition. + primary := c.primaryPartitionNode(partitionID) + if primary == nil { + return nil, errors.Errorf("translating index(%s) keys(%v) on partition(%d) - cannot find primary node", indexName, keys, partitionID) + } + + if c.Node.ID == primary.ID { + // The partition is local. + continue + } + + // Group the partition to be processed remotely. + keysByNode[primary] = append(keysByNode[primary], keys...) + + // Delete remote keys from the by-partition map so that it can be used for local translation. + delete(keysByPartition, partitionID) + } + + // Start translating keys remotely. + // On child calls, there are no remote results since we were only sent the keys that we own. + remoteResults := make(chan map[string]uint64, len(keysByNode)) + var g errgroup.Group + defer g.Wait() //nolint:errcheck + for node, keys := range keysByNode { + node, keys := node, keys + + g.Go(func() error { + translations, err := c.InternalClient.FindIndexKeysNode(ctx, &node.URI, indexName, keys...) + if err != nil { + return errors.Wrapf(err, "translating index(%s) keys(%v) on node %s", indexName, keys, node.ID) + } + + remoteResults <- translations + return nil + }) + } + + // Translate local keys. + translations := make(map[string]uint64) + for partitionID, keys := range keysByPartition { + // Handle cancellation. + select { + case <-done: + return nil, ctx.Err() + default: + } + + // Find the keys within the partition. + t, err := idx.TranslateStore(partitionID).FindKeys(keys...) + if err != nil { + return nil, errors.Wrapf(err, "translating index(%s) keys(%v) on partition(%d)", idx.Name(), keys, partitionID) + } + + // Merge the translations from this partition. + for key, id := range t { + translations[key] = id + } + } + + // Wait for remote key sets. + if err := g.Wait(); err != nil { + return nil, err + } + + // Merge the translations. + // All data should have been written to here while we waited. + // Closing the channel prevents the range from blocking. + close(remoteResults) + for t := range remoteResults { + for key, id := range t { + translations[key] = id + } + } + return translations, nil +} + +func (c *cluster) createIndexKeys(ctx context.Context, indexName string, keys ...string) (map[string]uint64, error) { + // Check for early cancellation. + done := ctx.Done() + select { + case <-done: + return nil, ctx.Err() + default: + } + + idx := c.holder.Index(indexName) + if idx == nil { + return nil, ErrIndexNotFound + } + + // Split keys by partition. + keysByPartition := make(map[int][]string, c.partitionN) + for _, key := range keys { + partitionID := c.keyPartition(indexName, key) + keysByPartition[partitionID] = append(keysByPartition[partitionID], key) + } + + // TODO: use local replicas to short-circuit network traffic + + // Group keys by node. + // Delete remote keys from the by-partition map so that it can be used for local translation. + keysByNode := make(map[*Node][]string) + for partitionID, keys := range keysByPartition { + // Find the primary node for this partition. + primary := c.primaryPartitionNode(partitionID) + if primary == nil { + return nil, errors.Errorf("translating index(%s) keys(%v) on partition(%d) - cannot find primary node", indexName, keys, partitionID) + } + + if c.Node.ID == primary.ID { + // The partition is local. + continue + } + + // Group the partition to be processed remotely. + keysByNode[primary] = append(keysByNode[primary], keys...) + delete(keysByPartition, partitionID) + } + + translateResults := make(chan map[string]uint64, len(keysByNode)+len(keysByPartition)) + var g errgroup.Group + defer g.Wait() //nolint:errcheck + + // Start translating keys remotely. + // On child calls, there are no remote results since we were only sent the keys that we own. + for node, keys := range keysByNode { + node, keys := node, keys + + g.Go(func() error { + translations, err := c.InternalClient.CreateIndexKeysNode(ctx, &node.URI, indexName, keys...) + if err != nil { + return errors.Wrapf(err, "translating index(%s) keys(%v) on node %s", indexName, keys, node.ID) + } + + translateResults <- translations + return nil + }) + } + + // Translate local keys. + // TODO: make this less horrible (why fsync why?????) + // This is kinda terrible because each goroutine does an fsync, thus locking up an entire OS thread. + // AHHHHHHHHHHHHHHHHHH + for partitionID, keys := range keysByPartition { + partitionID, keys := partitionID, keys + + g.Go(func() error { + // Handle cancellation. + select { + case <-done: + return ctx.Err() + default: + } + + translations, err := idx.TranslateStore(partitionID).CreateKeys(keys...) + if err != nil { + return errors.Wrapf(err, "translating index(%s) keys(%v) on partition(%d)", idx.Name(), keys, partitionID) + } + + translateResults <- translations + return nil + }) + } + + // Wait for remote key sets. + if err := g.Wait(); err != nil { + return nil, err + } + + // Merge the translations. + // All data should have been written to here while we waited. + // Closing the channel prevents the range from blocking. + translations := make(map[string]uint64, len(keys)) + close(translateResults) + for t := range translateResults { + for key, id := range t { + translations[key] = id + } + } + return translations, nil +} + func (c *cluster) translateIndexIDs(ctx context.Context, indexName string, ids []uint64) ([]string, error) { idSet := make(map[uint64]struct{}) for _, id := range ids { diff --git a/executor.go b/executor.go index 00de0b484..7f99970c9 100644 --- a/executor.go +++ b/executor.go @@ -218,64 +218,33 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar return resp, fmt.Errorf("profiling execution failed: %T is not tracing.Profile", prof) } } - resp.Results = make([]interface{}, len(q.Calls)) - for i, c := range q.Calls { - resp.Results[i] = emptyResult(c) + results, err := e.execute(ctx, index, q, shards, opt) + if err != nil { + return resp, err } - var columnAttrsRows []*Row - for i, c := range q.Calls { - // Translate query keys to ids, if necessary. - // No need to translate a remote call. - if !opt.Remote { - if err := e.translateCalls(ctx, index, []*pql.Call{c}); err != nil { - if errors.Cause(err) == ErrTranslatingKeyNotFound { - // No error - return empty result - continue - } - return resp, err - } else if err := validateQueryContext(ctx); err != nil { - return resp, err + var columnIDs []uint64 + if opt.ColumnAttrs { + // Consolidate all column ids across all calls. + for _, r := range results { + if bm, ok := r.(*Row); ok { + columnIDs = uint64Slice(columnIDs).merge(bm.Columns()) } } + } - results, err := e.execute(ctx, index, &pql.Query{Calls: []*pql.Call{c}}, shards, opt) + // Translate response objects from ids to keys, if necessary. + // No need to translate a remote call. + if !opt.Remote { + err = e.translateResults(ctx, index, idx, q.Calls, results) if err != nil { return resp, err - } else if err := validateQueryContext(ctx); err != nil { - return resp, err } - - if opt.ColumnAttrs { - if resultRow, ok := results[0].(*Row); ok { - columnAttrsRows = append(columnAttrsRows, resultRow) - } - } - - // Translate response objects from ids to keys, if necessary. - // No need to translate a remote call. - if !opt.Remote { - if err := e.translateResults(ctx, index, idx, []*pql.Call{c}, results); err != nil { - if errors.Cause(err) == ErrTranslatingKeyNotFound { - // No error - return empty result - continue - } - return resp, err - } else if err := validateQueryContext(ctx); err != nil { - return resp, err - } - } - - resp.Results[i] = results[0] } + resp.Results = results + // Fill column attributes if requested. if opt.ColumnAttrs { - // Consolidate all column ids across all calls. - var columnIDs []uint64 - for _, bm := range columnAttrsRows { - columnIDs = uint64Slice(columnIDs).merge(bm.Columns()) - } - // Retrieve column attributes across all calls. columnAttrSets, err := e.readColumnAttrSets(e.Holder.Index(index), columnIDs) if err != nil { @@ -455,6 +424,17 @@ func (e *executor) execute(ctx context.Context, index string, q *pql.Query, shar span, ctx := tracing.StartSpanFromContext(ctx, "Executor.execute") defer span.Finish() + // Apply translations if necessary. + var colTranslations map[string]map[string]uint64 // colID := colTranslations[index][key] + var rowTranslations map[string]map[string]map[string]uint64 // rowID := rowTranslations[index][field][key] + if !opt.Remote { + cols, rows, err := e.preTranslate(ctx, index, q.Calls...) + if err != nil { + return nil, err + } + colTranslations, rowTranslations = cols, rows + } + // Don't bother calculating shards for query types that don't require it. needsShards := needsShards(q.Calls) @@ -474,7 +454,7 @@ func (e *executor) execute(ctx context.Context, index string, q *pql.Query, shar // Optimize handling for bulk attribute insertion. if hasOnlySetRowAttrs(q.Calls) { - return e.executeBulkSetRowAttrs(ctx, index, q.Calls, opt) + return e.executeBulkSetRowAttrs(ctx, index, q.Calls, opt, colTranslations, rowTranslations) } // Execute each call serially. @@ -484,6 +464,20 @@ func (e *executor) execute(ctx context.Context, index string, q *pql.Query, shar return nil, err } + // Apply call translation. + if !opt.Remote { + translated, err := e.translateCall(call, index, colTranslations, rowTranslations) + if err != nil { + return nil, errors.Wrap(err, "translating call") + } + if translated == nil { + results = append(results, emptyResult(call)) + continue + } + + call = translated + } + // If you actually make a top-level Distinct call, you // want a SignedRow back. Otherwise, it's something else // that will be using it as a row, and we only care @@ -755,19 +749,9 @@ func (e *executor) executeFieldValueCall(ctx context.Context, index string, c *p return ValCount{}, ErrFieldNotFound } - var colID uint64 - if key, ok := colKey.(string); ok && idx.Keys() { - id, err := e.Cluster.translateIndexKey(ctx, index, key, false) - if err != nil { - return ValCount{}, errors.Wrap(err, "getting column id") - } - colID = id - } else { - id, ok, err := c.UintArg("column") - if !ok || err != nil { - return ValCount{}, errors.Wrap(err, "getting column argument") - } - colID = id + colID, ok, err := c.UintArg("column") + if !ok || err != nil { + return ValCount{}, errors.Wrap(err, "getting column argument") } shard := colID / ShardWidth @@ -3163,19 +3147,7 @@ func (e *executor) executeSetRow(ctx context.Context, indexName string, c *pql.C } field := e.Holder.Field(indexName, fieldName) if field == nil { - // Find index. - index := e.Holder.Index(indexName) - if index == nil { - return false, newNotFoundError(ErrIndexNotFound) - } - - // Create field. - field, err = index.CreateField(fieldName, OptFieldTypeSet(CacheTypeNone, 0)) - if err != nil { - // We wrap these because we want to indicate that it wasn't found, - // but also the problem we encountered trying to create it. - return false, newNotFoundError(errors.Wrap(err, "creating field")) - } + return false, errors.Wrapf(ErrFieldNotFound, "field %q", fieldName) } if field.Type() != FieldTypeSet { return false, fmt.Errorf("can't Store() on a %s field", field.Type()) @@ -3518,7 +3490,8 @@ func (e *executor) executeSetRowAttrs(ctx context.Context, index string, c *pql. } // executeBulkSetRowAttrs executes a set of SetRowAttrs() calls. -func (e *executor) executeBulkSetRowAttrs(ctx context.Context, index string, calls []*pql.Call, opt *execOptions) ([]interface{}, error) { + +func (e *executor) executeBulkSetRowAttrs(ctx context.Context, index string, calls []*pql.Call, opt *execOptions, colTranslations map[string]map[string]uint64, rowTranslations map[string]map[string]map[string]uint64) ([]interface{}, error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeBulkSetRowAttrs") defer span.Finish() @@ -3531,6 +3504,19 @@ func (e *executor) executeBulkSetRowAttrs(ctx context.Context, index string, cal } } + // Apply call translation. + if !opt.Remote { + translated, err := e.translateCall(c, index, colTranslations, rowTranslations) + if err != nil { + return nil, errors.Wrap(err, "translating call") + } + if translated == nil { + continue + } + + c = translated + } + field, ok := c.Args["_field"].(string) if !ok { return nil, errors.New("SetRowAttrs() field required") @@ -3931,287 +3917,667 @@ func (e *executor) mapperLocal(ctx context.Context, shards []uint64, mapFn mapFu } } -func (e *executor) translateCalls(ctx context.Context, defaultIndexName string, calls []*pql.Call) (err error) { - span, _ := tracing.StartSpanFromContext(ctx, "Executor.translateCalls") - defer span.Finish() - - // Generate a list of all used - keySets := make(map[string]map[string]struct{}) - for _, c := range calls { - if err := e.collectCallKeySets(ctx, defaultIndexName, c, keySets); err != nil { - return err +func (e *executor) preTranslate(ctx context.Context, index string, calls ...*pql.Call) (cols map[string]map[string]uint64, rows map[string]map[string]map[string]uint64, err error) { + // Collect all of the required keys. + collector := keyCollector{ + createCols: make(map[string][]string), + findCols: make(map[string][]string), + createRows: make(map[string]map[string][]string), + findRows: make(map[string]map[string][]string), + } + for _, call := range calls { + err := e.collectCallKeys(&collector, call, index) + if err != nil { + return nil, nil, err } } - // Perform a separate batch translation for each separate index used. - keyMaps := make(map[string]map[string]uint64) - for indexName, keySet := range keySets { - idx := e.Holder.indexes[indexName] + // Create keys. + // Both rows and columns need to be created first because of foreign index keys. + cols = make(map[string]map[string]uint64) + rows = make(map[string]map[string]map[string]uint64) + for index, keys := range collector.createCols { + translations, err := e.Cluster.createIndexKeys(ctx, index, keys...) + if err != nil { + return nil, nil, errors.Wrap(err, "creating query column keys") + } + cols[index] = translations + } + for index, fields := range collector.createRows { + idxRows := make(map[string]map[string]uint64) + idx := e.Holder.Index(index) if idx == nil { - return fmt.Errorf("cannot find index %q", indexName) + return nil, nil, errors.Wrapf(ErrIndexNotFound, "creating rows on index %q", index) } - - if !idx.Keys() || len(keySets) == 0 { - continue - } - if keyMaps[indexName], err = e.Cluster.translateIndexKeySet(ctx, indexName, keySet, true); err != nil { - return err - } - } - - // Translate calls. - for _, c := range calls { - if err := e.translateCall(ctx, defaultIndexName, c, keyMaps, true); err != nil { - return err - } - } - - return nil -} - -func (e *executor) collectCallKeySets(ctx context.Context, indexName string, c *pql.Call, m map[string]map[string]struct{}) error { - // Specifying an 'index' call overrides indexes on subsequent calls. - if s := c.CallIndex(); s != "" { - indexName = s - } - - if m[indexName] == nil { - m[indexName] = make(map[string]struct{}) - } - - // Collect key for this call. - colKey, rowKey, fieldName := c.TranslateInfo(columnLabel, rowLabel) - if c.Args[colKey] != nil && isString(c.Args[colKey]) { - if value := callArgString(c, colKey); value != "" { - m[indexName][value] = struct{}{} - } - } - - // Collect foreign index keys. - if fieldName != "" { - idx, exists := e.Holder.indexes[indexName] - if !exists { - return errors.Wrapf(ErrIndexNotFound, "%s", indexName) - } - if field := idx.Field(fieldName); field != nil && field.ForeignIndex() != "" { - foreignIndexName := field.ForeignIndex() - if m[foreignIndexName] == nil { - m[foreignIndexName] = make(map[string]struct{}) + for field, keys := range fields { + f := idx.Field(field) + if f == nil { + return nil, nil, errors.Wrapf(ErrFieldNotFound, "creating rows on field %q in index %q", field, index) } - - if c.Args[rowKey] != nil && isCondition(c.Args[rowKey]) { - cond := c.Args[rowKey].(*pql.Condition) - if isString(cond.Value) { - m[foreignIndexName][cond.Value.(string)] = struct{}{} - } - } else if value := callArgString(c, rowKey); value != "" { - m[foreignIndexName][value] = struct{}{} - } - } - } - - // Recursively collect argument calls. - for _, arg := range c.Args { - if arg, ok := arg.(*pql.Call); ok { - if err := e.collectCallKeySets(ctx, indexName, arg, m); err != nil { - return errors.Wrap(err, "collecting group by call index name") - } - } - } - - // Recursively collect child calls. - for _, child := range c.Children { - if err := e.collectCallKeySets(ctx, indexName, child, m); err != nil { - return err - } - } - return nil -} - -func (e *executor) translateCall(ctx context.Context, indexName string, c *pql.Call, keyMaps map[string]map[string]uint64, writable bool) (err error) { - // Specifying an 'index' arg applies to all nested calls. - if s := c.CallIndex(); s != "" { - indexName = s - } - keyMap := keyMaps[indexName] - - // Translate column key. - colKey, rowKey, fieldName := c.TranslateInfo(columnLabel, rowLabel) - idx, exists := e.Holder.indexes[indexName] - if !exists { - return errors.Wrapf(ErrIndexNotFound, "%s", indexName) - } - if idx.Keys() { - if c.Args[colKey] != nil && !isString(c.Args[colKey]) { - if !isValidID(c.Args[colKey]) { - return errors.Errorf("column value must be a string or non-negative integer, but got: %v of %[1]T", c.Args[colKey]) - } - } else if value := callArgString(c, colKey); value != "" { - c.Args[colKey] = keyMap[value] - } - } else { - if isString(c.Args[colKey]) { - return errors.New("string 'col' value not allowed unless index 'keys' option enabled") - } - } - - // Translate row key, if field is specified & key exists. - var field *Field - if fieldName != "" { - field = idx.Field(fieldName) - if field == nil { - // Instead of returning ErrFieldNotFound here, - // we just return, and don't attempt the translation. - // The assumption is that the non-existent field - // will raise an error downstream when it's used. - return nil - } - - // Bool field keys do not use the translator because there - // are only two possible values. Instead, they are handled - // directly. - if field.Type() == FieldTypeBool { - if c.Name == "Rows" { - // TranslateInfo for Rows returns "previous" as rowKey, - // so for bool fields we would get "missing bool argument" error - return nil - } - boolVal, err := callArgBool(c, rowKey) + translations, err := e.Cluster.createFieldKeys(ctx, f, keys...) if err != nil { - return errors.Wrapf(err, "getting bool key (%+v)", rowKey) + return nil, nil, errors.Wrap(err, "creating query row keys") } - rowID := falseRowID - if boolVal { - rowID = trueRowID - } - c.Args[rowKey] = rowID - } else if field.Keys() { - foreignIndexName := field.ForeignIndex() - if c.Args[rowKey] != nil && isCondition(c.Args[rowKey]) { - // In the case where a field has a foreign index with keys, - // allow `== "key"` or `!= "key"` to be used against the BSI - // field. - cond := c.Args[rowKey].(*pql.Condition) - if isString(cond.Value) { - switch cond.Op { - case pql.EQ, pql.NEQ: - var id uint64 - if foreignIndexName != "" { - id = keyMaps[foreignIndexName][cond.Value.(string)] - } else { - if id, err = e.Cluster.translateFieldKey(ctx, field, cond.Value.(string), writable); err != nil { - return errors.Wrapf(err, "translating field key: %s", cond.Value) - } - } + idxRows[field] = translations + } + rows[index] = idxRows + } - c.Args[rowKey] = &pql.Condition{ - Op: cond.Op, - Value: id, - } - default: - return errors.Errorf("conditional is not supported with string predicates: %s", cond.Op) - } - } - } else if c.Args[rowKey] != nil && !isString(c.Args[rowKey]) { - // allow passing row id directly (this can come in handy, but make sure it is a valid row id) - if !isValidID(c.Args[rowKey]) { - return errors.Errorf("row value must be a string or non-negative integer, but got: %v of %[1]T", c.Args[rowKey]) - } - } else if value := callArgString(c, rowKey); value != "" { - var id uint64 - if foreignIndexName != "" { - id = keyMaps[foreignIndexName][value] - } else { - if id, err = e.Cluster.translateFieldKey(ctx, field, value, writable); err != nil { - return errors.Wrapf(err, "translating field key: %s", value) - } - } - c.Args[rowKey] = id + // Find other keys. + for index, keys := range collector.findCols { + translations, err := e.Cluster.findIndexKeys(ctx, index, keys...) + if err != nil { + return nil, nil, errors.Wrap(err, "finding query column keys") + } + if prev := cols[index]; prev != nil { + for key, id := range translations { + prev[key] = id } } else { - if isString(c.Args[rowKey]) { - return errors.New("string 'row' value not allowed unless field 'keys' option enabled") + cols[index] = translations + } + } + for index, fields := range collector.findRows { + idxRows := rows[index] + if idxRows == nil { + idxRows = make(map[string]map[string]uint64) + rows[index] = idxRows + } + idx := e.Holder.Index(index) + if idx == nil { + return nil, nil, errors.Wrapf(ErrIndexNotFound, "finding rows on index %q", index) + } + for field, keys := range fields { + f := idx.Field(field) + if f == nil { + return nil, nil, errors.Wrapf(ErrFieldNotFound, "finding rows on field %q in index %q", field, index) + } + translations, err := e.Cluster.findFieldKeys(ctx, f, keys...) + if err != nil { + return nil, nil, errors.Wrap(err, "finding query row keys") + } + if prev := idxRows[field]; prev != nil { + for key, id := range translations { + prev[key] = id + } + } else { + idxRows[field] = translations + } + } + } + + return cols, rows, nil +} + +func (e *executor) collectCallKeys(dst *keyCollector, c *pql.Call, index string) error { + // Check for an overriding 'index' argument. + // This also applies to all child calls. + if callIndex := c.CallIndex(); callIndex != "" { + index = callIndex + } + + // Handle the field arg. + switch c.Name { + case "Set": + if field, err := c.FieldArg(); err == nil { + if arg, ok := c.Args[field].(string); ok { + dst.CreateRows(index, field, arg) + } + } + + case "Store": + if field, err := c.FieldArg(); err == nil { + idx := e.Holder.Index(index) + if idx == nil { + return errors.Wrapf(ErrIndexNotFound, "translating store field argument") + } + f := idx.Field(field) + if f == nil { + // Create the field. + // This is messy, because if a query leading up to the store fails, we will have created the field without executing the store. + var keyed bool + switch v := c.Args[field].(type) { + case string: + keyed = true + case uint64: + case int64: + if v < 0 { + return errors.Errorf("negative store row ID %d", v) + } + default: + return errors.Errorf("invalid store row identifier: %v of %T", v, v) + } + opts := []FieldOption{OptFieldTypeSet(CacheTypeNone, 0)} + if keyed { + opts = append(opts, OptFieldKeys()) + } + if _, err := idx.CreateField(field, opts...); err != nil { + // We wrap these because we want to indicate that it wasn't found, + // but also the problem we encountered trying to create it. + return newNotFoundError(errors.Wrapf(err, "creating field %q", field)) + } + } + if arg, ok := c.Args[field].(string); ok { + dst.CreateRows(index, field, arg) + } + } + + case "Clear", "Row", "Range", "ClearRow": + if field, err := c.FieldArg(); err == nil { + switch arg := c.Args[field].(type) { + case string: + dst.FindRows(index, field, arg) + case *pql.Condition: + // This is a workaround to allow `==` and `!=` to work on foreign index fields. + if key, ok := arg.Value.(string); ok { + switch arg.Op { + case pql.EQ, pql.NEQ: + dst.FindRows(index, field, key) + default: + return errors.Errorf("operator %v not defined on strings", arg.Op) + } + } + } + } + } + + // Handle _col. + if col, ok := c.Args["_col"].(string); ok { + switch c.Name { + case "Set", "SetColumnAttrs": + dst.CreateColumns(index, col) + default: + dst.FindColumns(index, col) + } + } + + // Handle _row. + if row, ok := c.Args["_row"].(string); ok { + // Find the field. + field, ok, err := c.StringArg("_field") + if err != nil { + return errors.Wrap(err, "finding field") + } + if !ok { + return errors.Wrap(ErrFieldNotFound, "finding field for _row argument") + } + + switch c.Name { + case "SetRowAttrs": + dst.CreateRows(index, field, row) + default: + dst.FindRows(index, field, row) + } + } + + // Handle queries that need a "column" argument. + switch c.Name { + case "Rows", "GroupBy", "FieldValue", "IncludesColumn": + if col, ok := c.Args["column"].(string); ok { + dst.FindColumns(index, col) + } + } + + // Handle special per-query arguments. + switch c.Name { + case "ConstRow": + // Translate the columns list. + if cols, ok := c.Args["columns"].([]interface{}); ok { + keys := make([]string, 0, len(cols)) + for _, v := range cols { + switch v := v.(type) { + case string: + keys = append(keys, v) + case uint64: + case int64: + default: + return errors.Errorf("invalid column identifier %v of type %T", c, c) + } + } + dst.FindColumns(index, keys...) + } + + case "Rows": + if prev, ok := c.Args["previous"].(string); ok { + // Find the field. + var field string + if f, ok, err := c.StringArg("_field"); err != nil { + return errors.Wrap(err, "finding field for Rows previous translation") + } else if ok { + field = f + } else if f, ok, err := c.StringArg("field"); err != nil { + return errors.Wrap(err, "finding field for Rows previous translation") + } else if ok { + field = f + } else { + return errors.New("missing field in Rows call") + } + + dst.FindRows(index, field, prev) + } + } + + // Collect keys from child calls. + for _, child := range c.Children { + err := e.collectCallKeys(dst, child, index) + if err != nil { + return err + } + } + + // Collect keys from argument calls. + for _, arg := range c.Args { + argCall, ok := arg.(*pql.Call) + if !ok { + continue + } + + err := e.collectCallKeys(dst, argCall, index) + if err != nil { + return err + } + } + + return nil +} + +type keyCollector struct { + createCols, findCols map[string][]string // map[index] -> column keys + createRows, findRows map[string]map[string][]string // map[index]map[field] -> row keys +} + +func (c *keyCollector) CreateColumns(index string, columns ...string) { + if len(columns) == 0 { + return + } + c.createCols[index] = append(c.createCols[index], columns...) +} + +func (c *keyCollector) FindColumns(index string, columns ...string) { + if len(columns) == 0 { + return + } + c.findCols[index] = append(c.findCols[index], columns...) +} + +func (c *keyCollector) CreateRows(index string, field string, columns ...string) { + if len(columns) == 0 { + return + } + idx := c.createRows[index] + if idx == nil { + idx = make(map[string][]string) + c.createRows[index] = idx + } + idx[field] = append(idx[field], columns...) +} + +func (c *keyCollector) FindRows(index string, field string, columns ...string) { + if len(columns) == 0 { + return + } + idx := c.findRows[index] + if idx == nil { + idx = make(map[string][]string) + c.findRows[index] = idx + } + idx[field] = append(idx[field], columns...) +} + +func fieldValidateValue(f *Field, val interface{}) error { + if val == nil { + return nil + } + + // Validate special types. + switch val := val.(type) { + case string: + if !f.Keys() { + return errors.Errorf("string value on an unkeyed field %q", f.Name()) + } + return nil + case *pql.Condition: + switch v := val.Value.(type) { + case nil: + case string: + case uint64: + case int64: + case float64: + case pql.Decimal: + case []interface{}: + for _, v := range v { + if err := fieldValidateValue(f, v); err != nil { + return err + } + } + return nil + default: + return errors.Errorf("invalid value %v in condition %q", v, val.String()) + } + return fieldValidateValue(f, val.Value) + } + + switch f.Type() { + case FieldTypeSet, FieldTypeMutex, FieldTypeTime: + switch v := val.(type) { + case uint64: + case int64: + if v < 0 { + return errors.Errorf("negative ID %d for set field %q", v, f.Name()) + } + default: + return errors.Errorf("invalid value %v for field %q of type %s", v, f.Name(), f.Type()) + } + case FieldTypeBool: + switch v := val.(type) { + case bool: + default: + return errors.Errorf("invalid value %v for bool field %q", v, f.Name()) + } + case FieldTypeInt: + switch v := val.(type) { + case uint64: + if v > 1<<63 { + return errors.Errorf("oversized integer %d for int field %q (range: -2^63 to 2^63-1)", v, f.Name()) + } + case int64: + default: + return errors.Errorf("invalid value %v for int field %q", v, f.Name()) + } + case FieldTypeDecimal: + switch v := val.(type) { + case uint64: + case int64: + case float64: + case pql.Decimal: + default: + return errors.Errorf("invalid value %v for decimal field %q", v, f.Name()) + } + default: + return errors.Errorf("unsupported type %s of field %q", f.Type(), f.Name()) + } + + return nil +} + +func (e *executor) translateCall(c *pql.Call, index string, columnKeys map[string]map[string]uint64, rowKeys map[string]map[string]map[string]uint64) (*pql.Call, error) { + // Check for an overriding 'index' argument. + // This also applies to all child calls. + if callIndex := c.CallIndex(); callIndex != "" { + index = callIndex + } + idx := e.Holder.Index(index) + if idx == nil { + return nil, errors.Wrapf(ErrIndexNotFound, "translating query on index %q", index) + } + + // Fetch the column keys list for this index. + indexCols, indexRows := columnKeys[index], rowKeys[index] + + // Handle the field arg. + switch c.Name { + case "Set", "Store": + if field, err := c.FieldArg(); err == nil { + f := e.Holder.Field(index, field) + if f == nil { + return nil, errors.Wrapf(ErrFieldNotFound, "validating value for field %q", field) + } + arg := c.Args[field] + if err := fieldValidateValue(f, arg); err != nil { + return nil, errors.Wrap(err, "validating store value") + } + switch arg := arg.(type) { + case string: + if translation, ok := indexRows[field][arg]; ok { + c.Args[field] = translation + } else { + return nil, errors.Wrapf(ErrTranslatingKeyNotFound, "destination key not found %q in %q in index %q", arg, field, index) + } + case bool: + if arg { + c.Args[field] = trueRowID + } else { + c.Args[field] = falseRowID + } + } + } + + case "Clear", "Row", "Range", "ClearRow": + if field, err := c.FieldArg(); err == nil { + f := e.Holder.Field(index, field) + if f == nil { + return nil, errors.Wrapf(ErrFieldNotFound, "validating value for field %q", field) + } + arg := c.Args[field] + if err := fieldValidateValue(f, arg); err != nil { + return nil, errors.Wrap(err, "validating field parameter value") + } + if c.Name == "Row" { + switch f.Type() { + case FieldTypeInt, FieldTypeDecimal: + if _, ok := arg.(*pql.Condition); !ok { + // This is workaround to support pql.ASSIGN ('=') as condition ('==') for int and decimal fields. + arg = &pql.Condition{ + Op: pql.EQ, + Value: arg, + } + c.Args[field] = arg + } + } + } + switch arg := arg.(type) { + case string: + if translation, ok := indexRows[field][arg]; ok { + c.Args[field] = translation + } else { + // Rewrite the call into a zero value call. + return e.callZero(c), nil + } + case bool: + if arg { + c.Args[field] = trueRowID + } else { + c.Args[field] = falseRowID + } + case *pql.Condition: + // This is a workaround to allow `==` and `!=` to work on foreign index fields. + if key, ok := arg.Value.(string); ok { + switch arg.Op { + case pql.EQ, pql.NEQ: + if translation, ok := indexRows[field][key]; ok { + arg.Value = translation + } else { + // Rewrite the call into a zero value call. + return e.callZero(c), nil + } + default: + return nil, errors.Errorf("operator %v not defined on strings", arg.Op) + } + } + } + } + } + + // Handle _col. + if col, ok := c.Args["_col"].(string); ok { + if !idx.Keys() { + return nil, errors.Wrapf(ErrTranslatingKeyNotFound, "translating column on unkeyed index %q", index) + } + if id, ok := indexCols[col]; ok { + c.Args["_col"] = id + } else { + switch c.Name { + case "Set", "SetColumnAttrs": + return nil, errors.Wrapf(ErrTranslatingKeyNotFound, "destination key not found %q in index %q", col, index) + default: + return e.callZero(c), nil + } + } + } + + // Handle _row. + if row, ok := c.Args["_row"]; ok { + // Find the field. + var field string + if f, ok, err := c.StringArg("_field"); err != nil { + return nil, errors.Wrap(err, "finding field") + } else if ok { + field = f + } else if f, ok, err := c.StringArg("field"); err != nil { + return nil, errors.Wrap(err, "finding field") + } else if ok { + field = f + } else { + return nil, errors.New("missing field") + } + + f := e.Holder.Field(index, field) + if f == nil { + return nil, errors.Wrapf(ErrFieldNotFound, "validating value for field %q", field) + } + if err := fieldValidateValue(f, row); err != nil { + return nil, errors.Wrap(err, "validating row value") + } + switch row := row.(type) { + case string: + if translation, ok := indexRows[field][row]; ok { + c.Args["_row"] = translation + } else { + switch c.Name { + case "SetRowAttrs": + return nil, errors.Errorf("row key missing in %q", c.String()) + default: + return e.callZero(c), nil + } + } + } + } + + // Handle queries that need a "column" argument. + switch c.Name { + case "Rows", "GroupBy", "FieldValue", "IncludesColumn": + if col, ok := c.Args["column"].(string); ok { + if translation, ok := indexCols[col]; ok { + c.Args["column"] = translation + } else { + // Rewrite the call into a zero value call. + return e.callZero(c), nil + } + } + } + + // Handle special per-query arguments. + switch c.Name { + case "ConstRow": + // Translate the columns list. + if cols, ok := c.Args["columns"].([]interface{}); ok { + out := make([]uint64, 0, len(cols)) + for _, v := range cols { + switch v := v.(type) { + case string: + if id, ok := indexCols[v]; ok { + out = append(out, id) + } + case uint64: + out = append(out, v) + case int64: + out = append(out, uint64(v)) + default: + return nil, errors.Errorf("invalid column identifier %v of type %T", c, c) + } + } + c.Args["columns"] = out + } + + case "Rows": + // Translate the previous row key. + if prev, ok := c.Args["previous"]; ok { + // Find the field. + var field string + if f, ok, err := c.StringArg("_field"); err != nil { + return nil, errors.Wrap(err, "finding field for Rows previous translation") + } else if ok { + field = f + } else if f, ok, err := c.StringArg("field"); err != nil { + return nil, errors.Wrap(err, "finding field for Rows previous translation") + } else if ok { + field = f + } else { + return nil, errors.New("missing field in Rows call") + } + + // Validate the type. + f := e.Holder.Field(index, field) + if f == nil { + return nil, errors.Wrapf(ErrFieldNotFound, "validating value for field %q", field) + } + if err := fieldValidateValue(f, prev); err != nil { + return nil, errors.Wrap(err, "validating prev value") + } + + switch prev := prev.(type) { + case string: + // Look up a translation for the previous row key. + if translation, ok := indexRows[field][prev]; ok { + c.Args["previous"] = translation + } else { + return nil, errors.Wrapf(ErrTranslatingKeyNotFound, "translating previous key %q from field %q in index %q in Rows call", prev, field, index) + } + case bool: + if prev { + c.Args["previous"] = trueRowID + } else { + c.Args["previous"] = falseRowID + } } } } // Translate child calls. - for _, child := range c.Children { - if err := e.translateCall(ctx, indexName, child, keyMaps, writable); err != nil { - return err + for i, child := range c.Children { + translated, err := e.translateCall(child, index, columnKeys, rowKeys) + if err != nil { + return nil, err } + c.Children[i] = translated } - // Translate call args. - for _, arg := range c.Args { - if arg, ok := arg.(*pql.Call); ok { - if err := e.translateCall(ctx, indexName, arg, keyMaps, writable); err != nil { - return errors.Wrap(err, "translating arg") - } - } - } - - // GroupBy-specific call translation. - if c.Name == "GroupBy" { - prev, ok := c.Args["previous"] + // Translate argument calls. + for k, arg := range c.Args { + argCall, ok := arg.(*pql.Call) if !ok { - return nil // nothing else to be translated - } - previous, ok := prev.([]interface{}) - if !ok { - return errors.Errorf("'previous' argument must be list, but got %T", prev) - } - if len(c.Children) != len(previous) { - return errors.Errorf("mismatched lengths for previous: %d and children: %d in %s", len(previous), len(c.Children), c) + continue } - fields := make([]*Field, len(c.Children)) - for i, child := range c.Children { - fieldname := callArgString(child, "_field") - field := idx.Field(fieldname) - if field == nil { - return errors.Wrapf(ErrFieldNotFound, "getting field '%s' from '%s'", fieldname, child) - } - fields[i] = field + translated, err := e.translateCall(argCall, index, columnKeys, rowKeys) + if err != nil { + return nil, err } - for i, field := range fields { - prev := previous[i] - if field.Keys() { - prevStr, ok := prev.(string) - if !ok { - return errors.New("prev value must be a string when field 'keys' option enabled") - } - // TODO: does this need to take field.ForeignIndex() into consideration? - id, err := e.Cluster.translateFieldKey(ctx, field, prevStr, writable) - if err != nil { - return errors.Wrapf(err, "translating field key: %s", prevStr) - } - previous[i] = id - } else { - if prevStr, ok := prev.(string); ok { - return errors.Errorf("got string row val '%s' in 'previous' for field %s which doesn't use string keys", prevStr, field.Name()) - } - } - } + c.Args[k] = translated } - // This is workaround to support pql.ASSIGN ('=') as condition ('==') for int and decimal fields - if c.Name == "Row" && field != nil && - (field.Type() == FieldTypeInt || field.Type() == FieldTypeDecimal) { - // re-write args as conditions for fieldName - for k, v := range c.Args { - if _, ok := v.(*pql.Condition); k == fieldName && !ok { - c.Args[k] = &pql.Condition{ - Op: pql.EQ, - Value: v, + return c, nil +} + +func (e *executor) callZero(c *pql.Call) *pql.Call { + switch c.Name { + case "Row", "Range": + if field, err := c.FieldArg(); err == nil { + if cond, ok := c.Args[field].(*pql.Condition); ok { + if cond.Op == pql.NEQ { + // Turn not nothing into everything. + return &pql.Call{Name: "All"} } - break } } - } - return nil + // Use an empty union as a placeholder. + return &pql.Call{Name: "Union"} + + default: + return nil + } } func (e *executor) translateResults(ctx context.Context, index string, idx *Index, calls []*pql.Call, results []interface{}) (err error) { @@ -4784,18 +5150,6 @@ func (vc *ValCount) floatLarger(other ValCount) ValCount { } } -func callArgBool(call *pql.Call, key string) (bool, error) { - value, ok := call.Args[key] - if !ok { - return false, errors.New("missing bool argument") - } - b, ok := value.(bool) - if !ok { - return false, fmt.Errorf("invalid bool argument type: %T", value) - } - return b, nil -} - func callArgString(call *pql.Call, key string) string { value, ok := call.Args[key] if !ok { @@ -4805,39 +5159,6 @@ func callArgString(call *pql.Call, key string) string { return s } -func isString(v interface{}) bool { - _, ok := v.(string) - return ok -} - -func isCondition(v interface{}) bool { - _, ok := v.(*pql.Condition) - return ok -} - -// isValidID returns whether v can be interpreted as a valid row or -// column ID. In short, is v a non-negative integer? I think the int64 -// and default cases are the only ones actually used since the PQL -// parser doesn't return any other integer types. -func isValidID(v interface{}) bool { - switch vt := v.(type) { - case uint, uint64, uint32, uint16, uint8: - return true - case int64: - return vt >= 0 - case int: - return vt >= 0 - case int32: - return vt >= 0 - case int16: - return vt >= 0 - case int8: - return vt >= 0 - default: - return false - } -} - // groupByIterator contains several slices. Each slice contains a number of // elements equal to the number of fields in the group by (the number of Rows // calls). diff --git a/executor_internal_test.go b/executor_internal_test.go index 534469864..622e7e7fa 100644 --- a/executor_internal_test.go +++ b/executor_internal_test.go @@ -20,121 +20,11 @@ import ( "fmt" "io/ioutil" "strconv" - "strings" "testing" "github.com/pilosa/pilosa/v2/pql" ) -func TestExecutor_TranslateGroupByCall(t *testing.T) { - holder := NewHolder(DefaultPartitionN) - - cluster := NewTestCluster(1) - - e := &executor{ - Holder: holder, - Cluster: cluster, - } - e.Holder.Path, _ = ioutil.TempDir(*TempDir, "") - err := e.Holder.Open() - if err != nil { - t.Fatalf("opening holder: %v", err) - } - - idx, err := e.Holder.CreateIndex("i", IndexOptions{}) - if err != nil { - t.Fatalf("creating index: %v", err) - } - - _, erra := idx.CreateField("ak", OptFieldKeys()) - _, errb := idx.CreateField("b") - _, errc := idx.CreateField("ck", OptFieldKeys()) - if erra != nil || errb != nil || errc != nil { - t.Fatalf("creating fields %v, %v, %v", erra, errb, errc) - } - - query, err := pql.ParseString(`GroupBy(Rows(ak), Rows(b), Rows(ck), previous=["la", 0, "ha"], having=Condition(count > 10))`) - if err != nil { - t.Fatalf("parsing query: %v", err) - } - c := query.Calls[0] - // this is writable call just for testing purpose - to test previous argument - // generally GroupBy calls are not writable and keys should already exist - err = e.translateCall(context.Background(), "i", c, make(map[string]map[string]uint64), true) - if err != nil { - t.Fatalf("translating call: %v", err) - } - if len(c.Args["previous"].([]interface{})) != 3 { - t.Fatalf("unexpected length for 'previous' arg %v", c.Args["previous"]) - } - for i, v := range c.Args["previous"].([]interface{}) { - if !isInt(v) { - t.Fatalf("expected all items in previous to be ints, but '%v' at index %d is %[1]T", v, i) - } - } - - if having, hok := c.Args["having"].(*pql.Call); !hok { - t.Fatal("expected having to be a call") - } else if cond, cok := having.Args["count"].(*pql.Condition); !cok { - t.Fatal("expected condition to be a count") - } else if cond.Op != pql.GT { - t.Fatal("expected condition op to be >") - } else { - val, ok := cond.Uint64Value() - if !ok || val != uint64(10) { - t.Fatal("expected condition val to be uint64(10)") - } - } - - errTests := []struct { - pql string - err string - }{ - { - pql: `GroupBy(Rows(notfound), previous=1)`, - err: "'previous' argument must be list", - }, - { - pql: `GroupBy(Rows(ak), previous=["la", 0])`, - err: "mismatched lengths", - }, - { - pql: `GroupBy(Rows(ak), previous=[1])`, - err: "prev value must be a string", - }, - { - pql: `GroupBy(Rows(notfound), previous=[1])`, - err: ErrFieldNotFound.Error(), - }, - // TODO: an unknown key will actually allocate an id. this is probably bad. - // { - // pql: `GroupBy(Rows(ak), previous=["zoop"])`, - // err: "translating row key '", - // }, - { - pql: `GroupBy(Rows(b), previous=["la"])`, - err: "which doesn't use string keys", - }, - } - - for i, test := range errTests { - t.Run(fmt.Sprintf("#%d_%s", i, test.err), func(t *testing.T) { - query, err := pql.ParseString(test.pql) - if err != nil { - t.Fatalf("parsing query: %v", err) - } - c := query.Calls[0] - err = e.translateCall(context.Background(), "i", c, make(map[string]map[string]uint64), false) - if err == nil { - t.Fatalf("expected error, but translated call is '%s", c) - } - if !strings.Contains(err.Error(), test.err) { - t.Fatalf("expected '%s', got '%v'", test.err, err) - } - }) - } -} - func TestExecutor_TranslateRowsOnBool(t *testing.T) { holder := NewHolder(DefaultPartitionN) defer holder.Close() @@ -183,7 +73,11 @@ func TestExecutor_TranslateRowsOnBool(t *testing.T) { } c := query.Calls[0] - err = e.translateCall(context.Background(), "i", c, make(map[string]map[string]uint64), true) + colTranslations, rowTranslations, err := e.preTranslate(context.Background(), "i", c) + if err != nil { + t.Fatalf("pre-translating call: %v", err) + } + _, err = e.translateCall(c, "i", colTranslations, rowTranslations) if err != nil { t.Fatalf("translating call: %v", err) } @@ -191,15 +85,6 @@ func TestExecutor_TranslateRowsOnBool(t *testing.T) { } } -func isInt(a interface{}) bool { - switch a.(type) { - case int, int64, uint, uint64: - return true - default: - return false - } -} - func TestFilterWithLimit(t *testing.T) { f := filterWithLimit(5) diff --git a/executor_test.go b/executor_test.go index df0c6a5de..39494c1ed 100644 --- a/executor_test.go +++ b/executor_test.go @@ -543,13 +543,13 @@ func TestExecutor_Execute_Set(t *testing.T) { }) t.Run("ErrInvalidColValueType", func(t *testing.T) { - if _, err := cmd.API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set("foo", f=1)`}); err == nil || errors.Cause(err).Error() != `string 'col' value not allowed unless index 'keys' option enabled` { + if _, err := cmd.API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set("foo", f=1)`}); err == nil || !hasCause(err, pilosa.ErrTranslatingKeyNotFound) || !strings.Contains(err.Error(), "unkeyed index") { t.Fatalf("The error is: '%v'", err) } }) t.Run("ErrInvalidRowValueType", func(t *testing.T) { - if _, err := cmd.API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set(2, f="bar")`}); err == nil || errors.Cause(err).Error() != `string 'row' value not allowed unless field 'keys' option enabled` { + if _, err := cmd.API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set(2, f="bar")`}); err == nil || !hasCause(err, pilosa.ErrTranslatingKeyNotFound) || !strings.Contains(err.Error(), "field is not keyed") { t.Fatal(err) } }) @@ -639,7 +639,7 @@ func TestExecutor_Execute_Set(t *testing.T) { if _, err := index.CreateField("f", pilosa.OptFieldTypeDefault(), pilosa.OptFieldKeys()); err != nil { t.Fatal(err) } - if _, err := cmd.API.Query(context.Background(), &pilosa.QueryRequest{Index: "inokey", Query: `Set(2, f=1.2)`}); err == nil || !strings.Contains(err.Error(), "row value must be a string or non-negative integer") { + if _, err := cmd.API.Query(context.Background(), &pilosa.QueryRequest{Index: "inokey", Query: `Set(2, f=1.2)`}); err == nil || !strings.Contains(err.Error(), "invalid value") { t.Fatal(err) } @@ -921,7 +921,7 @@ func TestExecutor_Execute_SetValue(t *testing.T) { } }) - t.Run("", func(t *testing.T) { + t.Run("Err", func(t *testing.T) { c := test.MustRunCluster(t, 1) defer c.Close() hldr := test.Holder{Holder: c[0].Server.Holder()} @@ -931,26 +931,39 @@ func TestExecutor_Execute_SetValue(t *testing.T) { t.Fatal(err) } - t.Run("ErrColumnBSIGroupRequired", func(t *testing.T) { - if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set(invalid_column_name=10, f=100)`}); err == nil || errors.Cause(err).Error() != `Set() column argument 'col' required` { + t.Run("ColumnBSIGroupRequired", func(t *testing.T) { + if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set(f=100)`}); err == nil || errors.Cause(err).Error() != `Set() column argument 'col' required` { t.Fatalf("unexpected error: %s", err) } }) - t.Run("ErrColumnBSIGroupValue", func(t *testing.T) { - if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set("bad_column", f=100)`}); err == nil || errors.Cause(err).Error() != `string 'col' value not allowed unless index 'keys' option enabled` { + t.Run("ColumnBSIGroupValue", func(t *testing.T) { + if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set("bad_column", f=100)`}); err == nil || !hasCause(err, pilosa.ErrTranslatingKeyNotFound) || !strings.Contains(err.Error(), "unkeyed index") { t.Fatalf("unexpected error: %s", err) } }) - t.Run("ErrInvalidBSIGroupValueType", func(t *testing.T) { - if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set(10, f="hello")`}); err == nil || errors.Cause(err).Error() != `string 'row' value not allowed unless field 'keys' option enabled` { + t.Run("InvalidBSIGroupValueType", func(t *testing.T) { + if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set(10, f="hello")`}); err == nil || !hasCause(err, pilosa.ErrTranslatingKeyNotFound) || !strings.Contains(err.Error(), "field is not keyed") { t.Fatalf("unexpected error: %s", err) } }) }) } +func hasCause(err, cause error) bool { + for err != cause { + innerErr := errors.Cause(err) + if innerErr == err { + // This is the innermost accessible error, and it does not have that cause. + return false + } + err = innerErr + } + + return true +} + // Ensure a SetRowAttrs() query can be executed. func TestExecutor_Execute_SetRowAttrs(t *testing.T) { c := test.MustRunCluster(t, 1) @@ -5907,3 +5920,91 @@ func TestTimelessClearRegression(t *testing.T) { t.Fatal("clear supposedly failed") } } + +func TestMissingKeyRegression(t *testing.T) { + c := test.MustRunCluster(t, 1) + defer c.Close() + + c.CreateField(t, "i", pilosa.IndexOptions{Keys: true, TrackExistence: true}, "f", pilosa.OptFieldKeys()) + + tests := []struct { + name string + query string + expected []interface{} + }{ + { + name: "RowGarbage", + query: `Row(f="garbage")`, + expected: []interface{}{[]string(nil)}, + }, + { + name: "Set", + query: `Set("a", f="example")`, + expected: []interface{}{true}, + }, + { + name: "Count", + query: `Count(Row(f="example"))`, + expected: []interface{}{uint64(1)}, + }, + { + name: "NotGarbage", + query: `Not(Row(f="garbage"))`, + expected: []interface{}{[]string{"a"}}, + }, + { + name: "DifferenceGarbage", + query: `Difference(All(), Row(f="garbage"))`, + expected: []interface{}{[]string{"a"}}, + }, + /*{ + // Key translation works here, but it seems the actual count query is processing stale data. + // Uncomment it when the bug has been fixed. + name: "SetAndCount", + query: `Set("a", f="example")` + "\n" + + `Count(Row(f="example"))`, + expected: []interface{}{true, uint64(1)}, + },*/ + { + name: "CountNothing", + query: `Count(Row(f="garbage"))`, + expected: []interface{}{uint64(0)}, + }, + { + name: "StoreInvertSelf", + query: `Store(Not(Row(f="xyzzy")), f="xyzzy")`, + expected: []interface{}{true}, + }, + { + name: "SetClear", + query: `Set("b", f="plugh")` + "\n" + + `Clear("b", f="plugh")`, + expected: []interface{}{true, true}, + }, + { + name: "ClearMix", + query: `Clear("a", f="garbage")` + "\n" + + `Clear("a", f="example")`, + expected: []interface{}{false, true}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + resp := c.Query(t, "i", tc.query) + if len(resp.Results) != len(tc.expected) { + t.Errorf("expected %d results but got %d", len(resp.Results), len(tc.expected)) + return + } + for i, r := range resp.Results { + if row, ok := r.(*pilosa.Row); ok { + r = row.Keys + } + expect := tc.expected[i] + if !reflect.DeepEqual(r, expect) { + t.Errorf("result %d differs: expected %v but got %v", i, expect, r) + } + } + }) + } +} diff --git a/http/client.go b/http/client.go index 986dd6afe..943604bf9 100644 --- a/http/client.go +++ b/http/client.go @@ -1247,6 +1247,208 @@ func (c *InternalClient) TranslateIDsNode(ctx context.Context, uri *pilosa.URI, return tkresp.Keys, nil } +func (c *InternalClient) FindIndexKeysNode(ctx context.Context, uri *pilosa.URI, index string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FindIndexKeysNode") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/keys/find", index)) + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("POST", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + +func (c *InternalClient) FindFieldKeysNode(ctx context.Context, uri *pilosa.URI, index string, field string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FindFieldKeysNode") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/field/%s/%s/keys/find", index, field)) + q := u.Query() + q.Add("remote", "true") + u.RawQuery = q.Encode() + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("POST", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + +func (c *InternalClient) CreateIndexKeysNode(ctx context.Context, uri *pilosa.URI, index string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.CreateIndexKeysNode") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/keys/create", index)) + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("POST", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + +func (c *InternalClient) CreateFieldKeysNode(ctx context.Context, uri *pilosa.URI, index string, field string, keys ...string) (transMap map[string]uint64, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.CreateFieldKeysNode") + defer span.Finish() + + // Create HTTP request. + u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/field/%s/%s/keys/create", index, field)) + q := u.Query() + q.Add("remote", "true") + u.RawQuery = q.Encode() + reqData, err := json.Marshal(keys) + if err != nil { + return nil, errors.Wrap(err, "marshalling request") + } + req, err := http.NewRequest("POST", u.String(), bytes.NewReader(reqData)) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + // Apply headers. + req.Header.Set("Content-Length", strconv.Itoa(len(reqData))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + + // Send the request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, errors.Wrap(err, "executing request") + } + defer func() { + cerr := resp.Body.Close() + if cerr != nil && err == nil { + err = errors.Wrap(cerr, "closing response body") + } + }() + + // Read the response body. + result, err := ioutil.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + + // Decode the translations. + transMap = make(map[string]uint64, len(keys)) + err = json.Unmarshal(result, &transMap) + if err != nil { + return nil, errors.Wrap(err, "json decoding") + } + + return transMap, nil +} + func (c *InternalClient) Transactions(ctx context.Context) (map[string]*pilosa.Transaction, error) { span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.Transactions") defer span.Finish() diff --git a/http/handler.go b/http/handler.go index d6162a5cb..7287efb0c 100644 --- a/http/handler.go +++ b/http/handler.go @@ -385,6 +385,11 @@ func newRouter(handler *Handler) *mux.Router { router.HandleFunc("/internal/nodes", handler.handleGetNodes).Methods("GET").Name("GetNodes") router.HandleFunc("/internal/shards/max", handler.handleGetShardsMax).Methods("GET").Name("GetShardsMax") // TODO: deprecate, but it's being used by the client + router.HandleFunc("/internal/translate/index/{index}/keys/find", handler.handleFindIndexKeys).Methods("POST").Name("FindIndexKeys") + router.HandleFunc("/internal/translate/index/{index}/keys/create", handler.handleCreateIndexKeys).Methods("POST").Name("CreateIndexKeys") + router.HandleFunc("/internal/translate/field/{index}/{field}/keys/find", handler.handleFindFieldKeys).Methods("POST").Name("FindFieldKeys") + router.HandleFunc("/internal/translate/field/{index}/{field}/keys/create", handler.handleCreateFieldKeys).Methods("POST").Name("CreateFieldKeys") + router.Use(handler.queryArgValidator) router.Use(handler.addQueryContext) router.Use(handler.extractTracing) @@ -2208,3 +2213,183 @@ func readBody(r *http.Request) ([]byte, error) { return buf.Bytes(), nil } + +func (h *Handler) handleFindIndexKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.FindIndexKeys(r.Context(), indexName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} + +func (h *Handler) handleFindFieldKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + fieldName, ok := mux.Vars(r)["field"] + if !ok { + http.Error(w, "field name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.FindFieldKeys(r.Context(), indexName, fieldName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} + +func (h *Handler) handleCreateIndexKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.CreateIndexKeys(r.Context(), indexName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} + +func (h *Handler) handleCreateFieldKeys(w http.ResponseWriter, r *http.Request) { + // Verify input and output types + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType) + return + } else if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "Not acceptable", http.StatusNotAcceptable) + return + } + + indexName, ok := mux.Vars(r)["index"] + if !ok { + http.Error(w, "index name is required", http.StatusBadRequest) + return + } + + fieldName, ok := mux.Vars(r)["field"] + if !ok { + http.Error(w, "field name is required", http.StatusBadRequest) + return + } + + bd, err := readBody(r) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var keys []string + err = json.Unmarshal(bd, &keys) + if err != nil { + http.Error(w, "failed to decode request", http.StatusBadRequest) + return + } + + translations, err := h.api.CreateFieldKeys(r.Context(), indexName, fieldName, keys...) + if err != nil { + http.Error(w, "translating keys", http.StatusBadRequest) + return + } + + err = json.NewEncoder(w).Encode(translations) + if err != nil { + http.Error(w, "encoding result", http.StatusBadRequest) + return + } +} diff --git a/mock/translator.go b/mock/translator.go index 8a28b504c..e7e88644f 100644 --- a/mock/translator.go +++ b/mock/translator.go @@ -21,8 +21,6 @@ import ( "github.com/pilosa/pilosa/v2" ) -var _ pilosa.TranslateStore = (*TranslateStore)(nil) - type TranslateStore struct { CloseFunc func() error MaxIDFunc func() (uint64, error) @@ -33,6 +31,8 @@ type TranslateStore struct { TranslateKeysFunc func(keys []string, writable bool) ([]uint64, error) TranslateIDFunc func(id uint64) (string, error) TranslateIDsFunc func(ids []uint64) ([]string, error) + FindKeysFunc func(keys ...string) (map[string]uint64, error) + CreateKeysFunc func(keys ...string) (map[string]uint64, error) ForceSetFunc func(id uint64, key string) error EntryReaderFunc func(ctx context.Context, offset uint64) (pilosa.TranslateEntryReader, error) } @@ -73,6 +73,14 @@ func (s *TranslateStore) TranslateIDs(ids []uint64) ([]string, error) { return s.TranslateIDsFunc(ids) } +func (s *TranslateStore) FindKeys(keys ...string) (map[string]uint64, error) { + return s.FindKeysFunc(keys...) +} + +func (s *TranslateStore) CreateKeys(keys ...string) (map[string]uint64, error) { + return s.CreateKeysFunc(keys...) +} + func (s *TranslateStore) ForceSet(id uint64, key string) error { return s.ForceSetFunc(id, key) } diff --git a/pql/ast.go b/pql/ast.go index e93a2aaf5..087a04120 100644 --- a/pql/ast.go +++ b/pql/ast.go @@ -662,6 +662,19 @@ func (c *Call) UintSliceArg(key string) ([]uint64, bool, error) { } } +func (c *Call) StringArg(key string) (string, bool, error) { + val, ok := c.Args[key] + if !ok { + return "", false, nil + } + switch tval := val.(type) { + case string: + return tval, true, nil + default: + return "", true, fmt.Errorf("unexpected type %T in StringArg, val %v", tval, tval) + } +} + // CallArg is for reading the value at key from call.Args as a Call. If the // key is not in Call.Args, the value of the returned value will be nil, and // the error will be nil. An error is returned if the value is not a Call. diff --git a/server/handler_test.go b/server/handler_test.go index e79248f2f..2d4b81fb5 100644 --- a/server/handler_test.go +++ b/server/handler_test.go @@ -603,7 +603,7 @@ func TestHandler_Endpoints(t *testing.T) { h.ServeHTTP(w, test.MustNewHTTPRequest("POST", "/index/i0/query", strings.NewReader(`Row(row=30)`))) if w.Code != gohttp.StatusBadRequest { t.Fatalf("unexpected status code: %d", w.Code) - } else if body := w.Body.String(); body != `{"error":"executing: map reduce: field not found"}`+"\n" { + } else if body := w.Body.String(); body != `{"error":"executing: translating call: validating value for field \"row\": field not found"}`+"\n" { t.Fatalf("unexpected body: %q", body) } }) @@ -620,7 +620,7 @@ func TestHandler_Endpoints(t *testing.T) { var resp pilosa.QueryResponse if err := cmd.API.Serializer.Unmarshal(w.Body.Bytes(), &resp); err != nil { t.Fatal(err) - } else if s := resp.Err.Error(); s != `executing: map reduce: field not found` { + } else if s := resp.Err.Error(); s != `executing: translating call: validating value for field "row": field not found` { t.Fatalf("unexpected error: %s", s) } }) diff --git a/stats/stats_test.go b/stats/stats_test.go index 8e9cb8de2..cb5d62232 100644 --- a/stats/stats_test.go +++ b/stats/stats_test.go @@ -197,7 +197,7 @@ func TestStatsCount_SetColumnAttrs(t *testing.T) { called := false idx := hldr.Holder.Index("d") if idx == nil { - t.Fatal("idex not found") + t.Fatal("index not found") } hldr.Holder.Stats = &MockStats{ diff --git a/translate.go b/translate.go index 4efa26129..618f6e5d8 100644 --- a/translate.go +++ b/translate.go @@ -50,7 +50,7 @@ var ( // 0xc2, 0xa0, // NO-BREAK SPACE // 0x00, // } -type TranslateStore interface { +type TranslateStore interface { // TODO: refactor this interface; readonly should be part of the type and replication should be an impl detail io.Closer // Returns the maximum ID set on the store. @@ -71,6 +71,15 @@ type TranslateStore interface { TranslateKey(key string, writable bool) (uint64, error) TranslateKeys(key []string, writable bool) ([]uint64, error) + // FindKeys looks up the ID for each key. + // Keys are not created if they do not exist. + // Missing keys are not considered errors, so the length of the result may be less than that of the input. + FindKeys(keys ...string) (map[string]uint64, error) + + // CreateKeys maps all keys to IDs, creating the IDs if they do not exist. + // If the translator is read-only, this will return an error. + CreateKeys(keys ...string) (map[string]uint64, error) + // Converts an integer ID to its associated string key. TranslateID(id uint64) (string, error) TranslateIDs(id []uint64) ([]string, error) @@ -335,6 +344,57 @@ func (s *InMemTranslateStore) TranslateKeys(keys []string, writable bool) (_ []u return ids, nil } +// FindKeys looks up the ID for each key. +// Keys are not created if they do not exist. +// Missing keys are not considered errors, so the length of the result may be less than that of the input. +func (s *InMemTranslateStore) FindKeys(keys ...string) (map[string]uint64, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + result := make(map[string]uint64, len(keys)) + for _, key := range keys { + id, ok := s.idsByKey[key] + if !ok { + // The key does not exist. + continue + } + + result[key] = id + } + + return result, nil +} + +// CreateKeys maps all keys to IDs, creating the IDs if they do not exist. +// If the translator is read-only, this will return an error. +func (s *InMemTranslateStore) CreateKeys(keys ...string) (map[string]uint64, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.readOnly { + return nil, ErrTranslateStoreReadOnly + } + + result := make(map[string]uint64, len(keys)) + for _, key := range keys { + id, ok := s.idsByKey[key] + if !ok { + // The key does not exist. + // Generate a new id and update db. + if s.field == "" { + id = GenerateNextPartitionedID(s.index, s.maxID, s.partitionID, s.partitionN) + } else { + id = s.maxID + 1 + } + s.set(id, key) + } + + result[key] = id + } + + return result, nil +} + func (s *InMemTranslateStore) translateKey(key string, writable bool) (_ uint64, err error) { id := s.idsByKey[key] if id != 0 { diff --git a/translator_test.go b/translator_test.go index 1229585c6..676fea363 100644 --- a/translator_test.go +++ b/translator_test.go @@ -17,11 +17,11 @@ package pilosa_test import ( "bytes" "context" - "errors" "fmt" "io" "reflect" "testing" + "time" "github.com/google/go-cmp/cmp" "github.com/pilosa/pilosa/v2" @@ -30,6 +30,8 @@ import ( "github.com/pilosa/pilosa/v2/mock" "github.com/pilosa/pilosa/v2/server" "github.com/pilosa/pilosa/v2/test" + "github.com/pkg/errors" + "golang.org/x/sync/errgroup" ) func TestInMemTranslateStore_TranslateKey(t *testing.T) { @@ -471,3 +473,210 @@ func TestTranslation_Coordinator(t *testing.T) { } }) } + +func TestTranslation_Cluster_CreateFind(t *testing.T) { + c := test.MustRunCluster(t, 3) + defer c.Close() + + c.CreateField(t, "i", pilosa.IndexOptions{Keys: true}, "f", pilosa.OptFieldKeys()) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + // Use the alphabet to test keys. + testKeys := make(map[string]struct{}) + for i := 'a'; i <= 'z'; i++ { + testKeys[string(i)] = struct{}{} + } + + t.Run("Index", func(t *testing.T) { + // Create all index keys, split across nodes. + { + parts := make([][]string, len(c)) + { + // Randomly partition the keys. + i := 0 + for k := range testKeys { + parts[i%len(c)] = append(parts[i%len(c)], k) + i++ + } + } + + // Create some keys on each node. + var g errgroup.Group + defer g.Wait() //nolint:errcheck + for i, keys := range parts { + i, keys := i, keys + g.Go(func() error { + _, err := c[i].API.CreateIndexKeys(ctx, "i", keys...) + return err + }) + } + if err := g.Wait(); err != nil { + t.Errorf("creating keys: %v", err) + return + } + } + + // Check that all index keys exist, and consistently map to the same IDs. + { + // Convert the keys to a list. + keyList := make([]string, 0, len(testKeys)) + for k := range testKeys { + keyList = append(keyList, k) + } + + // Obtain authoritative translations for the keys. + translations, err := c[0].API.FindIndexKeys(ctx, "i", keyList...) + if err != nil { + t.Errorf("obtaining authoritative translations: %v", err) + return + } + for _, k := range keyList { + if _, ok := translations[k]; !ok { + t.Errorf("key %q is missing", k) + } + } + + // Check that all nodes agree on these translations. + var g errgroup.Group + defer g.Wait() //nolint:errcheck + for i, n := range c { + i, api := i, n.API + g.Go(func() (err error) { + defer func() { err = errors.Wrapf(err, "translating on node %d", i) }() + localTranslations, err := api.FindIndexKeys(ctx, "i", keyList...) + if err != nil { + return errors.Wrap(err, "finding translations") + } + return compareTranslations(translations, localTranslations) + }) + } + if err := g.Wait(); err != nil { + t.Errorf("finding keys: %v", err) + return + } + + // Check that re-invoking create returns the original translations. + for i, n := range c { + i, api := i, n.API + g.Go(func() (err error) { + defer func() { err = errors.Wrapf(err, "translating on node %d", i) }() + localTranslations, err := api.CreateIndexKeys(ctx, "i", keyList...) + if err != nil { + return errors.Wrap(err, "finding translations") + } + return compareTranslations(translations, localTranslations) + }) + } + if err := g.Wait(); err != nil { + t.Errorf("checking re-create of keys: %v", err) + return + } + } + }) + t.Run("Field", func(t *testing.T) { + // Create all field keys, split across nodes. + { + parts := make([][]string, len(c)) + { + // Randomly partition the keys. + i := 0 + for k := range testKeys { + parts[i%len(c)] = append(parts[i%len(c)], k) + i++ + } + } + + // Create some keys on each node. + var g errgroup.Group + defer g.Wait() //nolint:errcheck + for i, keys := range parts { + i, keys := i, keys + g.Go(func() error { + _, err := c[i].API.CreateFieldKeys(ctx, "i", "f", keys...) + return err + }) + } + if err := g.Wait(); err != nil { + t.Errorf("creating keys: %v", err) + return + } + } + + // Check that all field keys exist, and consistently map to the same IDs. + { + // Convert the keys to a list. + keyList := make([]string, 0, len(testKeys)) + for k := range testKeys { + keyList = append(keyList, k) + } + + // Obtain authoritative translations for the keys. + translations, err := c[0].API.FindFieldKeys(ctx, "i", "f", keyList...) + if err != nil { + t.Errorf("obtaining authoritative translations: %v", err) + return + } + for _, k := range keyList { + if _, ok := translations[k]; !ok { + t.Errorf("key %q is missing", k) + } + } + + // Check that all nodes agree on these translations. + var g errgroup.Group + defer g.Wait() //nolint:errcheck + for i, n := range c { + i, api := i, n.API + g.Go(func() (err error) { + defer func() { err = errors.Wrapf(err, "translating on node %d", i) }() + localTranslations, err := api.FindFieldKeys(ctx, "i", "f", keyList...) + if err != nil { + return errors.Wrap(err, "finding translations") + } + return compareTranslations(translations, localTranslations) + }) + } + if err := g.Wait(); err != nil { + t.Errorf("finding keys: %v", err) + return + } + + // Check that re-invoking create returns the original translations. + for i, n := range c { + i, api := i, n.API + g.Go(func() (err error) { + defer func() { err = errors.Wrapf(err, "translating on node %d", i) }() + localTranslations, err := api.CreateFieldKeys(ctx, "i", "f", keyList...) + if err != nil { + return errors.Wrap(err, "finding translations") + } + return compareTranslations(translations, localTranslations) + }) + } + if err := g.Wait(); err != nil { + t.Errorf("checking re-create of keys: %v", err) + return + } + } + }) +} + +func compareTranslations(expected, got map[string]uint64) error { + for key, id := range got { + if realID, ok := expected[key]; !ok { + return errors.Errorf("unexpected key %q mapped to ID %d", key, id) + } else if id != realID { + return errors.Errorf("mismatched translation: expected %q:%d but got %q:%d", key, realID, key, id) + } + } + for key, realID := range expected { + if id, ok := got[key]; !ok { + return errors.Errorf("missing translation of key %q", key) + } else if id != realID { + return errors.Errorf("mismatched translation: expected %q:%d but got %q:%d", key, realID, key, id) + } + } + return nil +}