1package gocb
2
3import (
4	"encoding/json"
5	"log"
6
7	"github.com/opentracing/opentracing-go"
8	"gopkg.in/couchbase/gocbcore.v7"
9)
10
11type subDocResult struct {
12	path string
13	data []byte
14	err  error
15}
16
17// DocumentFragment represents multiple chunks of a full Document.
18type DocumentFragment struct {
19	cas      Cas
20	mt       MutationToken
21	contents []subDocResult
22	pathMap  map[string]int
23}
24
25// Cas returns the Cas of the Document
26func (frag *DocumentFragment) Cas() Cas {
27	return frag.cas
28}
29
30// MutationToken returns the MutationToken for the change represented by this DocumentFragment.
31func (frag *DocumentFragment) MutationToken() MutationToken {
32	return frag.mt
33}
34
35// ContentByIndex retrieves the value of the operation by its index. The index is the position of
36// the operation as it was added to the builder.
37func (frag *DocumentFragment) ContentByIndex(idx int, valuePtr interface{}) error {
38	res := frag.contents[idx]
39	if res.err != nil {
40		return res.err
41	}
42	if valuePtr == nil {
43		return nil
44	}
45
46	if valuePtr, ok := valuePtr.(*[]byte); ok {
47		*valuePtr = res.data
48		return nil
49	}
50
51	return json.Unmarshal(res.data, valuePtr)
52}
53
54// Content retrieves the value of the operation by its path. The path is the path provided
55// to the operation
56func (frag *DocumentFragment) Content(path string, valuePtr interface{}) error {
57	if frag.pathMap == nil {
58		frag.pathMap = make(map[string]int)
59		for i, v := range frag.contents {
60			frag.pathMap[v.path] = i
61		}
62	}
63	return frag.ContentByIndex(frag.pathMap[path], valuePtr)
64}
65
66// Exists checks whether the indicated path exists in this DocumentFragment and no
67// errors were returned from the server.
68func (frag *DocumentFragment) Exists(path string) bool {
69	err := frag.Content(path, nil)
70	return err == nil
71}
72
73// LookupInBuilder is a builder used to create a set of sub-document lookup operations.
74type LookupInBuilder struct {
75	bucket *Bucket
76	opName string
77	name   string
78	flags  gocbcore.SubdocDocFlag
79	ops    []gocbcore.SubDocOp
80}
81
82func (set *LookupInBuilder) execute(tracectx opentracing.SpanContext) (*DocumentFragment, error) {
83	return set.bucket.lookupIn(tracectx, set)
84}
85
86// Execute executes this set of lookup operations on the bucket.
87func (set *LookupInBuilder) Execute() (*DocumentFragment, error) {
88	return set.execute(nil)
89}
90
91// GetEx allows you to perform a sub-document Get operation with flags
92func (set *LookupInBuilder) GetEx(path string, flags SubdocFlag) *LookupInBuilder {
93	if path == "" {
94		op := gocbcore.SubDocOp{
95			Op:    gocbcore.SubDocOpGetDoc,
96			Flags: gocbcore.SubdocFlag(flags),
97		}
98		set.ops = append(set.ops, op)
99		return set
100	}
101
102	op := gocbcore.SubDocOp{
103		Op:    gocbcore.SubDocOpGet,
104		Path:  path,
105		Flags: gocbcore.SubdocFlag(flags),
106	}
107	set.ops = append(set.ops, op)
108	return set
109}
110
111// Get indicates a path to be retrieved from the document.  The value of the path
112// can later be retrieved (after .Execute()) using the Content or ContentByIndex
113// method. The path syntax follows N1QL's path syntax (e.g. `foo.bar.baz`).
114func (set *LookupInBuilder) Get(path string) *LookupInBuilder {
115	return set.GetEx(path, SubdocFlagNone)
116}
117
118// ExistsEx allows you to perform a sub-document Exists operation with flags
119func (set *LookupInBuilder) ExistsEx(path string, flags SubdocFlag) *LookupInBuilder {
120	op := gocbcore.SubDocOp{
121		Op:    gocbcore.SubDocOpExists,
122		Path:  path,
123		Flags: gocbcore.SubdocFlag(flags),
124	}
125	set.ops = append(set.ops, op)
126	return set
127}
128
129// Exists is similar to Get(), but does not actually retrieve the value from the server.
130// This may save bandwidth if you only need to check for the existence of a
131// path (without caring for its content). You can check the status of this
132// operation by using .Content (and ignoring the value) or .Exists()
133func (set *LookupInBuilder) Exists(path string) *LookupInBuilder {
134	return set.ExistsEx(path, SubdocFlagNone)
135}
136
137// GetCountEx allows you to perform a sub-document GetCount operation with flags
138func (set *LookupInBuilder) GetCountEx(path string, flags SubdocFlag) *LookupInBuilder {
139	op := gocbcore.SubDocOp{
140		Op:    gocbcore.SubDocOpGetCount,
141		Path:  path,
142		Flags: gocbcore.SubdocFlag(flags),
143	}
144	set.ops = append(set.ops, op)
145	return set
146}
147
148// GetCount allows you to retrieve the number of items in an array or keys within an
149// dictionary within an element of a document.
150func (set *LookupInBuilder) GetCount(path string) *LookupInBuilder {
151	return set.GetCountEx(path, SubdocFlagNone)
152}
153
154func (b *Bucket) lookupIn(tracectx opentracing.SpanContext, set *LookupInBuilder) (resOut *DocumentFragment, errOut error) {
155	if tracectx == nil {
156		lispan := b.startKvOpTrace(set.opName)
157		defer lispan.Finish()
158		tracectx = lispan.Context()
159	}
160
161	signal := make(chan bool, 1)
162	op, err := b.client.LookupInEx(gocbcore.LookupInOptions{
163		Key:          []byte(set.name),
164		Flags:        set.flags,
165		Ops:          set.ops,
166		TraceContext: tracectx,
167	}, func(res *gocbcore.LookupInResult, err error) {
168		errOut = err
169
170		if res != nil {
171			resSet := &DocumentFragment{}
172			resSet.contents = make([]subDocResult, len(res.Ops))
173			resSet.cas = Cas(res.Cas)
174
175			for i, opRes := range res.Ops {
176				resSet.contents[i].path = set.ops[i].Path
177				resSet.contents[i].err = opRes.Err
178				if opRes.Value != nil {
179					resSet.contents[i].data = append([]byte(nil), opRes.Value...)
180				}
181			}
182
183			resOut = resSet
184		}
185
186		signal <- true
187	})
188	if err != nil {
189		return nil, err
190	}
191
192	timeoutTmr := gocbcore.AcquireTimer(b.opTimeout)
193	select {
194	case <-signal:
195		gocbcore.ReleaseTimer(timeoutTmr, false)
196		return
197	case <-timeoutTmr.C:
198		gocbcore.ReleaseTimer(timeoutTmr, true)
199		if !op.Cancel() {
200			<-signal
201			return
202		}
203		return nil, ErrTimeout
204	}
205}
206
207func (b *Bucket) startLookupIn(opName string, key string, flags SubdocDocFlag) *LookupInBuilder {
208	return &LookupInBuilder{
209		bucket: b,
210		name:   key,
211		flags:  gocbcore.SubdocDocFlag(flags),
212		opName: opName,
213	}
214}
215
216// LookupInEx creates a sub-document lookup operation builder.
217func (b *Bucket) LookupInEx(key string, flags SubdocDocFlag) *LookupInBuilder {
218	return b.startLookupIn("LookupInEx", key, flags)
219}
220
221// LookupIn creates a sub-document lookup operation builder.
222func (b *Bucket) LookupIn(key string) *LookupInBuilder {
223	return b.startLookupIn("LookupIn", key, 0)
224}
225
226// MutateInBuilder is a builder used to create a set of sub-document mutation operations.
227type MutateInBuilder struct {
228	bucket    *Bucket
229	opName    string
230	name      string
231	flags     gocbcore.SubdocDocFlag
232	cas       gocbcore.Cas
233	expiry    uint32
234	ops       []gocbcore.SubDocOp
235	errs      MultiError
236	replicaTo uint
237	persistTo uint
238}
239
240func (set *MutateInBuilder) execute(tracectx opentracing.SpanContext) (*DocumentFragment, error) {
241	return set.bucket.mutateIn(tracectx, set)
242}
243
244// Execute executes this set of mutation operations on the bucket.
245func (set *MutateInBuilder) Execute() (*DocumentFragment, error) {
246	return set.execute(nil)
247}
248
249func (set *MutateInBuilder) marshalValue(value interface{}) []byte {
250	if value, ok := value.([]byte); ok {
251		return value
252	}
253
254	if value, ok := value.(*[]byte); ok {
255		return *value
256	}
257
258	bytes, err := json.Marshal(value)
259	if err != nil {
260		set.errs.add(err)
261		return nil
262	}
263	return bytes
264}
265
266// InsertEx allows you to perform a sub-document Insert operation with flags
267func (set *MutateInBuilder) InsertEx(path string, value interface{}, flags SubdocFlag) *MutateInBuilder {
268	if path == "" {
269		op := gocbcore.SubDocOp{
270			Op:    gocbcore.SubDocOpAddDoc,
271			Flags: gocbcore.SubdocFlag(flags),
272			Value: set.marshalValue(value),
273		}
274		set.ops = append(set.ops, op)
275		return set
276	}
277
278	op := gocbcore.SubDocOp{
279		Op:    gocbcore.SubDocOpDictAdd,
280		Path:  path,
281		Flags: gocbcore.SubdocFlag(flags),
282		Value: set.marshalValue(value),
283	}
284	set.ops = append(set.ops, op)
285	return set
286}
287
288// Insert adds an insert operation to this mutation operation set.
289func (set *MutateInBuilder) Insert(path string, value interface{}, createParents bool) *MutateInBuilder {
290	var flags SubdocFlag
291	if createParents {
292		flags |= SubdocFlagCreatePath
293	}
294
295	return set.InsertEx(path, value, flags)
296}
297
298// UpsertEx allows you to perform a sub-document Upsert operation with flags
299func (set *MutateInBuilder) UpsertEx(path string, value interface{}, flags SubdocFlag) *MutateInBuilder {
300	if path == "" {
301		op := gocbcore.SubDocOp{
302			Op:    gocbcore.SubDocOpSetDoc,
303			Flags: gocbcore.SubdocFlag(flags),
304			Value: set.marshalValue(value),
305		}
306		set.ops = append(set.ops, op)
307		return set
308	}
309
310	op := gocbcore.SubDocOp{
311		Op:    gocbcore.SubDocOpDictSet,
312		Path:  path,
313		Flags: gocbcore.SubdocFlag(flags),
314		Value: set.marshalValue(value),
315	}
316	set.ops = append(set.ops, op)
317	return set
318}
319
320// Upsert adds an upsert operation to this mutation operation set.
321func (set *MutateInBuilder) Upsert(path string, value interface{}, createParents bool) *MutateInBuilder {
322	var flags SubdocFlag
323	if createParents {
324		flags |= SubdocFlagCreatePath
325	}
326
327	return set.UpsertEx(path, value, flags)
328}
329
330// ReplaceEx allows you to perform a sub-document Replace operation with flags
331func (set *MutateInBuilder) ReplaceEx(path string, value interface{}, flags SubdocFlag) *MutateInBuilder {
332	op := gocbcore.SubDocOp{
333		Op:    gocbcore.SubDocOpReplace,
334		Path:  path,
335		Flags: gocbcore.SubdocFlag(flags),
336		Value: set.marshalValue(value),
337	}
338	set.ops = append(set.ops, op)
339	return set
340}
341
342// Replace adds an replace operation to this mutation operation set.
343func (set *MutateInBuilder) Replace(path string, value interface{}) *MutateInBuilder {
344	return set.ReplaceEx(path, value, SubdocFlagNone)
345}
346
347func (set *MutateInBuilder) marshalArrayMulti(in interface{}) (out []byte) {
348	out, err := json.Marshal(in)
349	if err != nil {
350		log.Panic(err)
351	}
352
353	// Assert first character is a '['
354	if len(out) < 2 || out[0] != '[' {
355		log.Panic("Not a JSON array")
356	}
357
358	out = out[1 : len(out)-1]
359	return
360}
361
362// RemoveEx allows you to perform a sub-document Remove operation with flags
363func (set *MutateInBuilder) RemoveEx(path string, flags SubdocFlag) *MutateInBuilder {
364	if path == "" {
365		op := gocbcore.SubDocOp{
366			Op:    gocbcore.SubDocOpDeleteDoc,
367			Flags: gocbcore.SubdocFlag(flags),
368		}
369		set.ops = append(set.ops, op)
370		return set
371	}
372
373	op := gocbcore.SubDocOp{
374		Op:    gocbcore.SubDocOpDelete,
375		Path:  path,
376		Flags: gocbcore.SubdocFlag(flags),
377	}
378	set.ops = append(set.ops, op)
379	return set
380}
381
382// Remove adds an remove operation to this mutation operation set.
383func (set *MutateInBuilder) Remove(path string) *MutateInBuilder {
384	return set.RemoveEx(path, SubdocFlagNone)
385}
386
387// ArrayPrependEx allows you to perform a sub-document ArrayPrepend operation with flags
388func (set *MutateInBuilder) ArrayPrependEx(path string, value interface{}, flags SubdocFlag) *MutateInBuilder {
389	return set.arrayPrependValue(path, set.marshalValue(value), flags)
390}
391
392// ArrayPrepend adds an element to the beginning (i.e. left) of an array
393func (set *MutateInBuilder) ArrayPrepend(path string, value interface{}, createParents bool) *MutateInBuilder {
394	var flags SubdocFlag
395	if createParents {
396		flags |= SubdocFlagCreatePath
397	}
398
399	return set.ArrayPrependEx(path, value, flags)
400}
401
402func (set *MutateInBuilder) arrayPrependValue(path string, bytes []byte, flags SubdocFlag) *MutateInBuilder {
403	op := gocbcore.SubDocOp{
404		Op:    gocbcore.SubDocOpArrayPushFirst,
405		Path:  path,
406		Flags: gocbcore.SubdocFlag(flags),
407		Value: bytes,
408	}
409	set.ops = append(set.ops, op)
410	return set
411}
412
413// ArrayAppendEx allows you to perform a sub-document ArrayAppend operation with flags
414func (set *MutateInBuilder) ArrayAppendEx(path string, value interface{}, flags SubdocFlag) *MutateInBuilder {
415	return set.arrayAppendValue(path, set.marshalValue(value), flags)
416}
417
418// ArrayAppend adds an element to the end (i.e. right) of an array
419func (set *MutateInBuilder) ArrayAppend(path string, value interface{}, createParents bool) *MutateInBuilder {
420	var flags SubdocFlag
421	if createParents {
422		flags |= SubdocFlagCreatePath
423	}
424
425	return set.ArrayAppendEx(path, value, flags)
426}
427
428func (set *MutateInBuilder) arrayAppendValue(path string, bytes []byte, flags SubdocFlag) *MutateInBuilder {
429	op := gocbcore.SubDocOp{
430		Op:    gocbcore.SubDocOpArrayPushLast,
431		Path:  path,
432		Flags: gocbcore.SubdocFlag(flags),
433		Value: bytes,
434	}
435	set.ops = append(set.ops, op)
436	return set
437}
438
439// ArrayInsertEx allows you to perform a sub-document ArrayInsert operation with flags
440func (set *MutateInBuilder) ArrayInsertEx(path string, value interface{}, flags SubdocFlag) *MutateInBuilder {
441	return set.arrayInsertValue(path, set.marshalValue(value), flags)
442}
443
444// ArrayInsert inserts an element at a given position within an array. The position should be
445// specified as part of the path, e.g. path.to.array[3]
446func (set *MutateInBuilder) ArrayInsert(path string, value interface{}) *MutateInBuilder {
447	return set.ArrayInsertEx(path, value, SubdocFlagNone)
448}
449
450func (set *MutateInBuilder) arrayInsertValue(path string, bytes []byte, flags SubdocFlag) *MutateInBuilder {
451	op := gocbcore.SubDocOp{
452		Op:    gocbcore.SubDocOpArrayInsert,
453		Path:  path,
454		Flags: gocbcore.SubdocFlag(flags),
455		Value: bytes,
456	}
457	set.ops = append(set.ops, op)
458	return set
459}
460
461// ArrayAppendMultiEx allows you to perform a sub-document ArrayAppendMulti operation with flags
462func (set *MutateInBuilder) ArrayAppendMultiEx(path string, values interface{}, flags SubdocFlag) *MutateInBuilder {
463	return set.arrayAppendValue(path, set.marshalArrayMulti(values), flags)
464}
465
466// ArrayAppendMulti adds multiple values as elements to an array.
467// `values` must be an array type
468// ArrayAppendMulti("path", []int{1,2,3,4}, true) =>
469//   "path" [..., 1,2,3,4]
470//
471// This is a more efficient version (at both the network and server levels)
472// of doing
473// ArrayAppend("path", 1, true).ArrayAppend("path", 2, true).ArrayAppend("path", 3, true)
474//
475// See ArrayAppend() for more information
476func (set *MutateInBuilder) ArrayAppendMulti(path string, values interface{}, createParents bool) *MutateInBuilder {
477	var flags SubdocFlag
478	if createParents {
479		flags |= SubdocFlagCreatePath
480	}
481
482	return set.ArrayAppendMultiEx(path, values, flags)
483}
484
485// ArrayPrependMultiEx allows you to perform a sub-document ArrayPrependMulti operation with flags
486func (set *MutateInBuilder) ArrayPrependMultiEx(path string, values interface{}, flags SubdocFlag) *MutateInBuilder {
487	return set.arrayPrependValue(path, set.marshalArrayMulti(values), flags)
488}
489
490// ArrayPrependMulti adds multiple values at the beginning of an array.
491// See ArrayAppendMulti for more information about multiple element operations
492// and ArrayPrepend for the semantics of this operation
493func (set *MutateInBuilder) ArrayPrependMulti(path string, values interface{}, createParents bool) *MutateInBuilder {
494	var flags SubdocFlag
495	if createParents {
496		flags |= SubdocFlagCreatePath
497	}
498
499	return set.ArrayPrependMultiEx(path, values, flags)
500}
501
502// ArrayInsertMultiEx allows you to perform a sub-document ArrayInsertMulti operation with flags
503func (set *MutateInBuilder) ArrayInsertMultiEx(path string, values interface{}, flags SubdocFlag) *MutateInBuilder {
504	return set.arrayInsertValue(path, set.marshalArrayMulti(values), flags)
505}
506
507// ArrayInsertMulti inserts multiple elements at a specified position within the
508// array. See ArrayAppendMulti for more information about multiple element
509// operations, and ArrayInsert for more information about array insertion operations
510func (set *MutateInBuilder) ArrayInsertMulti(path string, values interface{}) *MutateInBuilder {
511	return set.ArrayInsertMultiEx(path, values, SubdocFlagNone)
512}
513
514// ArrayAddUniqueEx allows you to perform a sub-document ArrayAddUnique operation with flags
515func (set *MutateInBuilder) ArrayAddUniqueEx(path string, value interface{}, flags SubdocFlag) *MutateInBuilder {
516	op := gocbcore.SubDocOp{
517		Op:    gocbcore.SubDocOpArrayAddUnique,
518		Path:  path,
519		Flags: gocbcore.SubdocFlag(flags),
520		Value: set.marshalValue(value),
521	}
522	set.ops = append(set.ops, op)
523	return set
524}
525
526// ArrayAddUnique adds an dictionary add unique operation to this mutation operation set.
527func (set *MutateInBuilder) ArrayAddUnique(path string, value interface{}, createParents bool) *MutateInBuilder {
528	var flags SubdocFlag
529	if createParents {
530		flags |= SubdocFlagCreatePath
531	}
532
533	return set.ArrayAddUniqueEx(path, value, flags)
534}
535
536// CounterEx allows you to perform a sub-document Counter operation with flags
537func (set *MutateInBuilder) CounterEx(path string, delta int64, flags SubdocFlag) *MutateInBuilder {
538	op := gocbcore.SubDocOp{
539		Op:    gocbcore.SubDocOpCounter,
540		Path:  path,
541		Flags: gocbcore.SubdocFlag(flags),
542		Value: set.marshalValue(delta),
543	}
544	set.ops = append(set.ops, op)
545	return set
546}
547
548// Counter adds an counter operation to this mutation operation set.
549func (set *MutateInBuilder) Counter(path string, delta int64, createParents bool) *MutateInBuilder {
550	var flags SubdocFlag
551	if createParents {
552		flags |= SubdocFlagCreatePath
553	}
554
555	return set.CounterEx(path, delta, flags)
556}
557
558func (b *Bucket) mutateIn(tracectx opentracing.SpanContext, set *MutateInBuilder) (resOut *DocumentFragment, errOut error) {
559	// Perform the base operation
560	res, err := b.mutateInBase(tracectx, set)
561	if err != nil {
562		return res, err
563	}
564
565	// Skip durability if there was none set
566	if set.replicaTo == 0 && set.persistTo == 0 {
567		return res, err
568	}
569
570	// Attempt to satisfy the durability requirements
571	return res, b.durability(tracectx, set.name, res.cas, res.mt, set.replicaTo, set.persistTo, false)
572}
573
574func (b *Bucket) mutateInBase(tracectx opentracing.SpanContext, set *MutateInBuilder) (resOut *DocumentFragment, errOut error) {
575	if tracectx == nil {
576		mispan := b.startKvOpTrace(set.opName)
577		defer mispan.Finish()
578		tracectx = mispan.Context()
579	}
580
581	errOut = set.errs.get()
582	if errOut != nil {
583		return
584	}
585
586	signal := make(chan bool, 1)
587	op, err := b.client.MutateInEx(gocbcore.MutateInOptions{
588		Key:          []byte(set.name),
589		Flags:        set.flags,
590		Cas:          set.cas,
591		Expiry:       set.expiry,
592		Ops:          set.ops,
593		TraceContext: tracectx,
594	}, func(res *gocbcore.MutateInResult, err error) {
595		errOut = err
596
597		if res != nil {
598			resSet := &DocumentFragment{
599				cas: Cas(res.Cas),
600				mt:  MutationToken{res.MutationToken, b},
601			}
602			resSet.contents = make([]subDocResult, len(res.Ops))
603
604			for i, opRes := range res.Ops {
605				resSet.contents[i].path = set.ops[i].Path
606				resSet.contents[i].err = opRes.Err
607				if opRes.Value != nil {
608					resSet.contents[i].data = append([]byte(nil), opRes.Value...)
609				}
610			}
611
612			resOut = resSet
613		}
614
615		signal <- true
616	})
617	if err != nil {
618		return nil, err
619	}
620
621	timeoutTmr := gocbcore.AcquireTimer(b.opTimeout)
622	select {
623	case <-signal:
624		gocbcore.ReleaseTimer(timeoutTmr, false)
625		return
626	case <-timeoutTmr.C:
627		gocbcore.ReleaseTimer(timeoutTmr, true)
628		if !op.Cancel() {
629			<-signal
630			return
631		}
632		return nil, ErrTimeout
633	}
634}
635
636func (b *Bucket) startMutateIn(opName string, key string, flags SubdocDocFlag, cas Cas, expiry uint32, replicaTo, persistTo uint) *MutateInBuilder {
637	return &MutateInBuilder{
638		bucket:    b,
639		opName:    opName,
640		name:      key,
641		flags:     gocbcore.SubdocDocFlag(flags),
642		cas:       gocbcore.Cas(cas),
643		expiry:    expiry,
644		replicaTo: replicaTo,
645		persistTo: persistTo,
646	}
647}
648
649// MutateInEx creates a sub-document mutation operation builder.
650func (b *Bucket) MutateInEx(key string, flags SubdocDocFlag, cas Cas, expiry uint32) *MutateInBuilder {
651	return b.startMutateIn("MutateInEx", key, flags, cas, expiry, 0, 0)
652}
653
654// MutateInExDura creates a sub-document mutation operation builder with durability.
655func (b *Bucket) MutateInExDura(key string, flags SubdocDocFlag, cas Cas, expiry uint32, replicaTo, persistTo uint) *MutateInBuilder {
656	return b.startMutateIn("MutateInExDura", key, flags, cas, expiry, replicaTo, persistTo)
657}
658
659// MutateIn creates a sub-document mutation operation builder.
660func (b *Bucket) MutateIn(key string, cas Cas, expiry uint32) *MutateInBuilder {
661	return b.startMutateIn("MutateIn", key, 0, cas, expiry, 0, 0)
662}
663