From 08c3c5b3778024e02e6c1186635dd923ef70e2ab Mon Sep 17 00:00:00 2001
From: Nikita Tokarchuk <tokarchuk+bitbucket@nevisq.com>
Date: Mon, 13 Jul 2020 04:08:19 +0200
Subject: [PATCH] Add callback mechanism and implement on-decode callback

---
 mongox/database/callbacks.go    | 19 +++++++++++++++++++
 mongox/database/loadarray.go    | 18 ++++++++++++++----
 mongox/database/loadone.go      | 12 +++++++++++-
 mongox/database/loadstream.go   |  2 +-
 mongox/database/streamloader.go | 10 ++++++++++
 mongox/query/callbacks.go       |  7 +++++++
 mongox/query/compose.go         | 12 ++++++++++++
 mongox/query/query.go           |  6 ++++++
 8 files changed, 80 insertions(+), 6 deletions(-)
 create mode 100644 mongox/database/callbacks.go
 create mode 100644 mongox/query/callbacks.go

diff --git a/mongox/database/callbacks.go b/mongox/database/callbacks.go
new file mode 100644
index 0000000..e062f5a
--- /dev/null
+++ b/mongox/database/callbacks.go
@@ -0,0 +1,19 @@
+package database
+
+import (
+	"context"
+
+	"github.com/mainnika/mongox-go-driver/v2/mongox/query"
+)
+
+func onDecode(ctx context.Context, iter interface{}, callbacks ...query.OnDecode) (err error) {
+
+	for _, cb := range callbacks {
+		err = cb(ctx, iter)
+		if err != nil {
+			return
+		}
+	}
+
+	return
+}
diff --git a/mongox/database/loadarray.go b/mongox/database/loadarray.go
index 0c702e3..b6eb7c9 100644
--- a/mongox/database/loadarray.go
+++ b/mongox/database/loadarray.go
@@ -49,14 +49,18 @@ func (d *Database) LoadArray(target interface{}, filters ...interface{}) (err er
 	}
 
 	for i = 0; result.Next(d.Context()); {
+
+		var elem interface{}
+
 		if targetSliceV.Len() == i {
-			elem := reflect.New(targetSliceElemT.Elem())
-			err = result.Decode(elem.Interface())
+			value := reflect.New(targetSliceElemT.Elem())
+			err = result.Decode(value.Interface())
+			elem = value.Interface()
 			if err == nil {
-				targetSliceV = reflect.Append(targetSliceV, elem)
+				targetSliceV = reflect.Append(targetSliceV, value)
 			}
 		} else {
-			elem := targetSliceV.Index(i).Interface()
+			elem = targetSliceV.Index(i).Interface()
 			base.Reset(elem)
 			err = result.Decode(elem)
 		}
@@ -65,6 +69,12 @@ func (d *Database) LoadArray(target interface{}, filters ...interface{}) (err er
 			return
 		}
 
+		err = onDecode(d.ctx, elem, composed.OnDecode()...)
+		if err != nil {
+			_ = result.Close(d.Context())
+			return
+		}
+
 		i++
 	}
 
diff --git a/mongox/database/loadone.go b/mongox/database/loadone.go
index fdb4e0a..37cef08 100644
--- a/mongox/database/loadone.go
+++ b/mongox/database/loadone.go
@@ -35,5 +35,15 @@ func (d *Database) LoadOne(target interface{}, filters ...interface{}) (err erro
 
 	base.Reset(target)
 
-	return result.Decode(target)
+	err = result.Decode(target)
+	if err != nil {
+		return
+	}
+
+	err = onDecode(d.ctx, target, composed.OnDecode()...)
+	if err != nil {
+		return
+	}
+
+	return
 }
diff --git a/mongox/database/loadstream.go b/mongox/database/loadstream.go
index 1e3339b..ee55585 100644
--- a/mongox/database/loadstream.go
+++ b/mongox/database/loadstream.go
@@ -25,7 +25,7 @@ func (d *Database) LoadStream(target interface{}, filters ...interface{}) (loade
 		return
 	}
 
-	loader = &StreamLoader{cur: cursor, ctx: d.Context(), target: target}
+	loader = &StreamLoader{cur: cursor, ctx: d.Context(), target: target, query: composed}
 
 	return
 }
diff --git a/mongox/database/streamloader.go b/mongox/database/streamloader.go
index ce1c0a3..2efe11c 100644
--- a/mongox/database/streamloader.go
+++ b/mongox/database/streamloader.go
@@ -5,11 +5,13 @@ import (
 
 	"github.com/mainnika/mongox-go-driver/v2/mongox"
 	"github.com/mainnika/mongox-go-driver/v2/mongox/base"
+	"github.com/mainnika/mongox-go-driver/v2/mongox/query"
 )
 
 // StreamLoader is a controller for a database cursor
 type StreamLoader struct {
 	cur    *mongox.Cursor
+	query  *query.Query
 	ctx    context.Context
 	target interface{}
 }
@@ -36,6 +38,14 @@ func (l *StreamLoader) Decode() (err error) {
 	base.Reset(l.target)
 
 	err = l.cur.Decode(l.target)
+	if err != nil {
+		return
+	}
+
+	err = onDecode(l.ctx, l.target, l.query.OnDecode()...)
+	if err != nil {
+		return
+	}
 
 	return
 }
diff --git a/mongox/query/callbacks.go b/mongox/query/callbacks.go
new file mode 100644
index 0000000..26adeee
--- /dev/null
+++ b/mongox/query/callbacks.go
@@ -0,0 +1,7 @@
+package query
+
+import (
+	"context"
+)
+
+type OnDecode func(ctx context.Context, iter interface{}) (err error)
diff --git a/mongox/query/compose.go b/mongox/query/compose.go
index 3929e44..d8b64b4 100644
--- a/mongox/query/compose.go
+++ b/mongox/query/compose.go
@@ -38,6 +38,7 @@ func Push(q *Query, f interface{}) (ok bool) {
 	ok = ok || applySkip(q, f)
 	ok = ok || applyProtection(q, f)
 	ok = ok || applyPreloader(q, f)
+	ok = ok || applyCallbacks(q, f)
 
 	return ok
 }
@@ -123,3 +124,14 @@ func applyPreloader(q *Query, f interface{}) (ok bool) {
 
 	return false
 }
+
+func applyCallbacks(q *Query, f interface{}) (ok bool) {
+
+	switch cb := f.(type) {
+	case OnDecode:
+		q.ondecode = append(q.ondecode, cb)
+		ok = true
+	}
+
+	return
+}
diff --git a/mongox/query/query.go b/mongox/query/query.go
index 7e4dafe..643b6e0 100644
--- a/mongox/query/query.go
+++ b/mongox/query/query.go
@@ -11,6 +11,7 @@ type Query struct {
 	sorter    Sorter
 	skipper   Skipper
 	preloader Preloader
+	ondecode  []OnDecode
 }
 
 // And function pushes the elem query to the $and array of the query
@@ -75,6 +76,11 @@ func (q *Query) Preloader() (ok bool, preloads []string) {
 	return
 }
 
+// OnDecode callback is called after the mongo decode function
+func (q *Query) OnDecode() (callbacks []OnDecode) {
+	return q.ondecode
+}
+
 // Empty checks the query for any content
 func (q *Query) Empty() (isEmpty bool) {
 	return len(q.m) == 0