Merge pull request #94 from nitrictech/feature/document-key-return

Document Key Return

pkg/adapters/grpc/document_grpc.go

@@ -118,6 +118,7 @@ func documentToWire(doc *sdk.Document) (*pb.Document, error) {
 
 	return &pb.Document{
 		Content: valStruct,
+		Key:     keyToWire(doc.Key),
 	}, nil
 }
 
@@ -136,6 +137,28 @@ func keyFromWire(key *pb.Key) *sdk.Key {
 	return sdkKey
 }
 
+// keyToWire - translates an SDK key to a gRPC key
+func keyToWire(key *sdk.Key) *pb.Key {
+	return &pb.Key{
+		Id:         key.Id,
+		Collection: collectionToWire(key.Collection),
+	}
+}
+
+// collectionToWire - translates a SDK collection to a gRPC collection
+func collectionToWire(col *sdk.Collection) *pb.Collection {
+	if col.Parent != nil {
+		return &pb.Collection{
+			Name:   col.Name,
+			Parent: keyToWire(col.Parent),
+		}
+	} else {
+		return &pb.Collection{
+			Name: col.Name,
+		}
+	}
+}
+
 // collectionFromWire - returns an Membrane SDK Document Collection from the protobuf wire representation
 // recursively calls keyFromWire if the collection is a sub-collection under another key
 func collectionFromWire(coll *pb.Collection) *sdk.Collection {

pkg/plugins/document/boltdb/boltdb.go

@@ -16,13 +16,14 @@ package boltdb_service
 
 import (
 	"fmt"
-	"github.com/nitric-dev/membrane/pkg/plugins/document"
-	"github.com/nitric-dev/membrane/pkg/utils"
 	"os"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/nitric-dev/membrane/pkg/plugins/document"
+	"github.com/nitric-dev/membrane/pkg/utils"
+
 	"github.com/Knetic/govaluate"
 	"github.com/asdine/storm"
 	"github.com/asdine/storm/q"
@@ -73,7 +74,7 @@ func (s *BoltDocService) Get(key *sdk.Key) (*sdk.Document, error) {
 		return nil, err
 	}
 
-	return toSdkDoc(doc), nil
+	return toSdkDoc(key.Collection, doc), nil
 }
 
 func (s *BoltDocService) Set(key *sdk.Key, content map[string]interface{}) error {
@@ -205,28 +206,19 @@ func (s *BoltDocService) Query(collection *sdk.Collection, expressions []sdk.Que
 	documents := make([]sdk.Document, 0)
 	scanCount := 0
 	for _, doc := range docs {
+		scanCount += 1
 
 		if filterExp != nil {
-			eval, err := filterExp.Evaluate(doc.Value)
-			if err != nil {
-				//return nil, err
-				//fmt.Printf("failed to evaluate query expression for document. Details: %v\n\tExpression: %v\n\tDocument: %v\n", err, filterExp.String(), doc.Value)
+			include, err := filterExp.Evaluate(doc.Value)
+			if err != nil || !(include.(bool)) {
 				// TODO: determine if skipping failed evaluations is always appropriate.
-				// Treat a failed eval as a mismatch, since it's usually a datatype mismatch or a missing key/prop on the doc, which is essentially a failed match.
+				// 	errors are usually a datatype mismatch or a missing key/prop on the doc, which is essentially a failed match.
+				// Treat a failed or false eval as a mismatch
 				continue
 			}
-			include := eval.(bool)
-			if include {
-				sdkDoc := toSdkDoc(doc)
-				documents = append(documents, *sdkDoc)
-			}
-
-		} else {
-			sdkDoc := toSdkDoc(doc)
-			documents = append(documents, *sdkDoc)
 		}
-
-		scanCount += 1
+		sdkDoc := toSdkDoc(collection, doc)
+		documents = append(documents, *sdkDoc)
 
 		// Break if greater than fetch limit
 		if limit > 0 && len(documents) == limit {
@@ -281,7 +273,6 @@ func (s *BoltDocService) createdDb(coll sdk.Collection) (*storm.DB, error) {
 }
 
 func createDoc(key *sdk.Key) BoltDoc {
-
 	parentKey := key.Collection.Parent
 
 	// Top Level Collection
@@ -301,8 +292,32 @@ func createDoc(key *sdk.Key) BoltDoc {
 	}
 }
 
-func toSdkDoc(doc BoltDoc) *sdk.Document {
+func toSdkDoc(col *sdk.Collection, doc BoltDoc) *sdk.Document {
+	keys := strings.Split(doc.Id, "_")
+
+	// Translate the boltdb Id into a nitric document key Id
+	var id string
+	var c *sdk.Collection
+	if len(keys) > 1 {
+		// sub document
+		id = keys[len(keys)-1]
+		c = &sdk.Collection{
+			Name: col.Name,
+			Parent: &sdk.Key{
+				Collection: col.Parent.Collection,
+				Id:         keys[0],
+			},
+		}
+	} else {
+		id = doc.Id
+		c = col
+	}
+
 	return &sdk.Document{
 		Content: doc.Value,
+		Key: &sdk.Key{
+			Collection: c,
+			Id:         id,
+		},
 	}
 }

pkg/plugins/document/dynamodb/dynamodb.go

@@ -17,6 +17,7 @@ package dynamodb_service
 import (
 	"fmt"
 	"sort"
+	"strings"
 
 	"github.com/nitric-dev/membrane/pkg/plugins/document"
 	"github.com/nitric-dev/membrane/pkg/utils"
@@ -74,6 +75,7 @@ func (s *DynamoDocService) Get(key *sdk.Key) (*sdk.Document, error) {
 	delete(itemMap, ATTRIB_SK)
 
 	return &sdk.Document{
+		Key:     key,
 		Content: itemMap,
 	}, nil
 }
@@ -333,7 +335,7 @@ func (s *DynamoDocService) performQuery(
 		return fmt.Errorf("error performing query %v: %v", input, err)
 	}
 
-	return marshalQueryResult(resp.Items, resp.LastEvaluatedKey, limit, queryResult)
+	return marshalQueryResult(collection, resp.Items, resp.LastEvaluatedKey, limit, queryResult)
 }
 
 func (s *DynamoDocService) performScan(
@@ -407,10 +409,11 @@ func (s *DynamoDocService) performScan(
 		return fmt.Errorf("error performing scan %v: %v", input, err)
 	}
 
-	return marshalQueryResult(resp.Items, resp.LastEvaluatedKey, limit, queryResult)
+	return marshalQueryResult(collection, resp.Items, resp.LastEvaluatedKey, limit, queryResult)
 }
 
 func marshalQueryResult(
+	collection *sdk.Collection,
 	items []map[string]*dynamodb.AttributeValue,
 	lastEvaluatedKey map[string]*dynamodb.AttributeValue,
 	limit int,
@@ -425,10 +428,40 @@ func marshalQueryResult(
 
 	// Strip keys & append results
 	for _, m := range valueMaps {
+		// Retrieve the original ID on the result
+		var id string
+		var c *sdk.Collection
+		if collection.Parent == nil {
+			// We know this is a root document so its key will be located in PK
+			pk, _ := m[ATTRIB_PK].(string)
+			id = pk
+			c = collection
+		} else {
+			// We know this is a child document so its key will be located in the SK
+			pk, _ := m[ATTRIB_PK].(string)
+			sk, _ := m[ATTRIB_SK].(string)
+			idStr := strings.Split(sk, "#")
+			id = idStr[len(idStr)-1]
+			c = &sdk.Collection{
+				Name: collection.Name,
+				Parent: &sdk.Key{
+					Collection: &sdk.Collection{
+						Name: collection.Parent.Collection.Name,
+					},
+					Id: pk,
+				},
+			}
+		}
+
+		// Split out sort key value
 		delete(m, ATTRIB_PK)
 		delete(m, ATTRIB_SK)
 
 		sdkDoc := sdk.Document{
+			Key: &sdk.Key{
+				Collection: c,
+				Id:         id,
+			},
 			Content: m,
 		}
 		queryResult.Documents = append(queryResult.Documents, sdkDoc)

pkg/plugins/document/firestore/firestore.go

@@ -17,9 +17,10 @@ package firestore_service
 import (
 	"context"
 	"fmt"
-	"github.com/nitric-dev/membrane/pkg/plugins/document"
 	"strings"
 
+	"github.com/nitric-dev/membrane/pkg/plugins/document"
+
 	"cloud.google.com/go/firestore"
 	"cloud.google.com/go/pubsub"
 	"github.com/nitric-dev/membrane/pkg/sdk"
@@ -49,6 +50,7 @@ func (s *FirestoreDocService) Get(key *sdk.Key) (*sdk.Document, error) {
 	}
 
 	return &sdk.Document{
+		Key:     key,
 		Content: value.Data(),
 	}, nil
 }
@@ -155,7 +157,24 @@ func (s *FirestoreDocService) Query(collection *sdk.Collection, expressions []sd
 		if err != nil {
 			return nil, fmt.Errorf("error querying value: %v", err)
 		}
-		sdkDoc := sdk.Document{Content: docSnp.Data()}
+		sdkDoc := sdk.Document{
+			Content: docSnp.Data(),
+			Key: &sdk.Key{
+				Collection: collection,
+				Id:         docSnp.Ref.ID,
+			},
+		}
+
+		if p := docSnp.Ref.Parent.Parent; p != nil {
+			sdkDoc.Key.Collection = &sdk.Collection{
+				Name: collection.Name,
+				Parent: &sdk.Key{
+					Collection: collection.Parent.Collection,
+					Id:         p.ID,
+				},
+			}
+		}
+
 		queryResult.Documents = append(queryResult.Documents, sdkDoc)
 
 		// If query limit configured determine continue tokens
@@ -220,14 +239,12 @@ func (s *FirestoreDocService) getQueryRoot(collection *sdk.Collection) firestore
 
 	if parentKey == nil {
 		return s.client.Collection(collection.Name).Offset(0)
-
 	} else {
 		if parentKey.Id != "" {
 			return s.client.Collection(parentKey.Collection.Name).
 				Doc(parentKey.Id).
 				Collection(collection.Name).
 				Offset(0)
-
 		} else {
 			// Note there is a risk of subcollection name collison
 			// TODO: future YAML validation could help mitigate this

pkg/sdk/document.go

@@ -37,6 +37,7 @@ func (k *Key) String() string {
 }
 
 type Document struct {
+	Key     *Key
 	Content map[string]interface{}
 }