blob: f2209a0f36209b85a62143f3f4eccd6766801315 [file] [log] [blame]
Don Newton379ae252019-04-01 12:17:06 -04001// Copyright (C) MongoDB, Inc. 2017-present.
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may
4// not use this file except in compliance with the License. You may obtain
5// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
6
7package bsonx
8
9import (
10 "bytes"
11 "errors"
12 "fmt"
13
14 "github.com/mongodb/mongo-go-driver/bson/bsontype"
15 "github.com/mongodb/mongo-go-driver/x/bsonx/bsoncore"
16)
17
18// ErrNilDocument indicates that an operation was attempted on a nil *bson.Document.
19var ErrNilDocument = errors.New("document is nil")
20
21// KeyNotFound is an error type returned from the Lookup methods on Document. This type contains
22// information about which key was not found and if it was actually not found or if a component of
23// the key except the last was not a document nor array.
24type KeyNotFound struct {
25 Key []string // The keys that were searched for.
26 Depth uint // Which key either was not found or was an incorrect type.
27 Type bsontype.Type // The type of the key that was found but was an incorrect type.
28}
29
30func (knf KeyNotFound) Error() string {
31 depth := knf.Depth
32 if depth >= uint(len(knf.Key)) {
33 depth = uint(len(knf.Key)) - 1
34 }
35
36 if len(knf.Key) == 0 {
37 return "no keys were provided for lookup"
38 }
39
40 if knf.Type != bsontype.Type(0) {
41 return fmt.Sprintf(`key "%s" was found but was not valid to traverse BSON type %s`, knf.Key[depth], knf.Type)
42 }
43
44 return fmt.Sprintf(`key "%s" was not found`, knf.Key[depth])
45}
46
47// Doc is a type safe, concise BSON document representation.
48type Doc []Elem
49
50// ReadDoc will create a Document using the provided slice of bytes. If the
51// slice of bytes is not a valid BSON document, this method will return an error.
52func ReadDoc(b []byte) (Doc, error) {
53 doc := make(Doc, 0)
54 err := doc.UnmarshalBSON(b)
55 if err != nil {
56 return nil, err
57 }
58 return doc, nil
59}
60
61// Copy makes a shallow copy of this document.
62func (d Doc) Copy() Doc {
63 d2 := make(Doc, len(d))
64 copy(d2, d)
65 return d2
66}
67
68// Append adds an element to the end of the document, creating it from the key and value provided.
69func (d Doc) Append(key string, val Val) Doc {
70 return append(d, Elem{Key: key, Value: val})
71}
72
73// Prepend adds an element to the beginning of the document, creating it from the key and value provided.
74func (d Doc) Prepend(key string, val Val) Doc {
75 // TODO: should we just modify d itself instead of doing an alloc here?
76 return append(Doc{{Key: key, Value: val}}, d...)
77}
78
79// Set replaces an element of a document. If an element with a matching key is
80// found, the element will be replaced with the one provided. If the document
81// does not have an element with that key, the element is appended to the
82// document instead.
83func (d Doc) Set(key string, val Val) Doc {
84 idx := d.IndexOf(key)
85 if idx == -1 {
86 return append(d, Elem{Key: key, Value: val})
87 }
88 d[idx] = Elem{Key: key, Value: val}
89 return d
90}
91
92// IndexOf returns the index of the first element with a key of key, or -1 if no element with a key
93// was found.
94func (d Doc) IndexOf(key string) int {
95 for i, e := range d {
96 if e.Key == key {
97 return i
98 }
99 }
100 return -1
101}
102
103// Delete removes the element with key if it exists and returns the updated Doc.
104func (d Doc) Delete(key string) Doc {
105 idx := d.IndexOf(key)
106 if idx == -1 {
107 return d
108 }
109 return append(d[:idx], d[idx+1:]...)
110}
111
112// Lookup searches the document and potentially subdocuments or arrays for the
113// provided key. Each key provided to this method represents a layer of depth.
114//
115// This method will return an empty Value if they key does not exist. To know if they key actually
116// exists, use LookupErr.
117func (d Doc) Lookup(key ...string) Val {
118 val, _ := d.LookupErr(key...)
119 return val
120}
121
122// LookupErr searches the document and potentially subdocuments or arrays for the
123// provided key. Each key provided to this method represents a layer of depth.
124func (d Doc) LookupErr(key ...string) (Val, error) {
125 elem, err := d.LookupElementErr(key...)
126 return elem.Value, err
127}
128
129// LookupElement searches the document and potentially subdocuments or arrays for the
130// provided key. Each key provided to this method represents a layer of depth.
131//
132// This method will return an empty Element if they key does not exist. To know if they key actually
133// exists, use LookupElementErr.
134func (d Doc) LookupElement(key ...string) Elem {
135 elem, _ := d.LookupElementErr(key...)
136 return elem
137}
138
139// LookupElementErr searches the document and potentially subdocuments for the
140// provided key. Each key provided to this method represents a layer of depth.
141func (d Doc) LookupElementErr(key ...string) (Elem, error) {
142 // KeyNotFound operates by being created where the error happens and then the depth is
143 // incremented by 1 as each function unwinds. Whenever this function returns, it also assigns
144 // the Key slice to the key slice it has. This ensures that the proper depth is identified and
145 // the proper keys.
146 if len(key) == 0 {
147 return Elem{}, KeyNotFound{Key: key}
148 }
149
150 var elem Elem
151 var err error
152 idx := d.IndexOf(key[0])
153 if idx == -1 {
154 return Elem{}, KeyNotFound{Key: key}
155 }
156
157 elem = d[idx]
158 if len(key) == 1 {
159 return elem, nil
160 }
161
162 switch elem.Value.Type() {
163 case bsontype.EmbeddedDocument:
164 switch tt := elem.Value.primitive.(type) {
165 case Doc:
166 elem, err = tt.LookupElementErr(key[1:]...)
167 case MDoc:
168 elem, err = tt.LookupElementErr(key[1:]...)
169 }
170 default:
171 return Elem{}, KeyNotFound{Type: elem.Value.Type()}
172 }
173 switch tt := err.(type) {
174 case KeyNotFound:
175 tt.Depth++
176 tt.Key = key
177 return Elem{}, tt
178 case nil:
179 return elem, nil
180 default:
181 return Elem{}, err // We can't actually hit this.
182 }
183}
184
185// MarshalBSONValue implements the bsoncodec.ValueMarshaler interface.
186//
187// This method will never return an error.
188func (d Doc) MarshalBSONValue() (bsontype.Type, []byte, error) {
189 if d == nil {
190 // TODO: Should we do this?
191 return bsontype.Null, nil, nil
192 }
193 data, _ := d.MarshalBSON()
194 return bsontype.EmbeddedDocument, data, nil
195}
196
197// MarshalBSON implements the Marshaler interface.
198//
199// This method will never return an error.
200func (d Doc) MarshalBSON() ([]byte, error) { return d.AppendMarshalBSON(nil) }
201
202// AppendMarshalBSON marshals Doc to BSON bytes, appending to dst.
203//
204// This method will never return an error.
205func (d Doc) AppendMarshalBSON(dst []byte) ([]byte, error) {
206 idx, dst := bsoncore.ReserveLength(dst)
207 for _, elem := range d {
208 t, data, _ := elem.Value.MarshalBSONValue() // Value.MarshalBSONValue never returns an error.
209 dst = append(dst, byte(t))
210 dst = append(dst, elem.Key...)
211 dst = append(dst, 0x00)
212 dst = append(dst, data...)
213 }
214 dst = append(dst, 0x00)
215 dst = bsoncore.UpdateLength(dst, idx, int32(len(dst[idx:])))
216 return dst, nil
217}
218
219// UnmarshalBSON implements the Unmarshaler interface.
220func (d *Doc) UnmarshalBSON(b []byte) error {
221 if d == nil {
222 return ErrNilDocument
223 }
224
225 if err := bsoncore.Document(b).Validate(); err != nil {
226 return err
227 }
228
229 elems, err := bsoncore.Document(b).Elements()
230 if err != nil {
231 return err
232 }
233 var val Val
234 for _, elem := range elems {
235 rawv := elem.Value()
236 err = val.UnmarshalBSONValue(rawv.Type, rawv.Data)
237 if err != nil {
238 return err
239 }
240 *d = d.Append(elem.Key(), val)
241 }
242 return nil
243}
244
245// UnmarshalBSONValue implements the bson.ValueUnmarshaler interface.
246func (d *Doc) UnmarshalBSONValue(t bsontype.Type, data []byte) error {
247 if t != bsontype.EmbeddedDocument {
248 return fmt.Errorf("cannot unmarshal %s into a bsonx.Doc", t)
249 }
250 return d.UnmarshalBSON(data)
251}
252
253// Equal compares this document to another, returning true if they are equal.
254func (d Doc) Equal(id IDoc) bool {
255 switch tt := id.(type) {
256 case Doc:
257 d2 := tt
258 if len(d) != len(d2) {
259 return false
260 }
261 for idx := range d {
262 if !d[idx].Equal(d2[idx]) {
263 return false
264 }
265 }
266 case MDoc:
267 unique := make(map[string]struct{}, 0)
268 for _, elem := range d {
269 unique[elem.Key] = struct{}{}
270 val, ok := tt[elem.Key]
271 if !ok {
272 return false
273 }
274 if !val.Equal(elem.Value) {
275 return false
276 }
277 }
278 if len(unique) != len(tt) {
279 return false
280 }
281 case nil:
282 return d == nil
283 default:
284 return false
285 }
286
287 return true
288}
289
290// String implements the fmt.Stringer interface.
291func (d Doc) String() string {
292 var buf bytes.Buffer
293 buf.Write([]byte("bson.Document{"))
294 for idx, elem := range d {
295 if idx > 0 {
296 buf.Write([]byte(", "))
297 }
298 fmt.Fprintf(&buf, "%v", elem)
299 }
300 buf.WriteByte('}')
301
302 return buf.String()
303}
304
305func (Doc) idoc() {}