From c62e9bcabea2b3ed5e3a8bc1602cd38f8ad477f9 Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Sun, 15 Jan 2017 18:03:12 +0800
Subject: [PATCH] Query Row, Rows inside RowQuery callbacks

---
 callback_row_query.go | 30 ++++++++++++++++++++++++++++++
 scope.go              | 14 ++++++++++----
 2 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 callback_row_query.go

diff --git a/callback_row_query.go b/callback_row_query.go
new file mode 100644
index 00000000..c2ff4a08
--- /dev/null
+++ b/callback_row_query.go
@@ -0,0 +1,30 @@
+package gorm
+
+import "database/sql"
+
+// Define callbacks for row query
+func init() {
+	DefaultCallback.RowQuery().Register("gorm:row_query", rowQueryCallback)
+}
+
+type RowQueryResult struct {
+	Row *sql.Row
+}
+
+type RowsQueryResult struct {
+	Rows  *sql.Rows
+	Error error
+}
+
+// queryCallback used to query data from database
+func rowQueryCallback(scope *Scope) {
+	if result, ok := scope.InstanceGet("row_query_result"); ok {
+		scope.prepareQuerySQL()
+
+		if rowResult, ok := result.(*RowQueryResult); ok {
+			rowResult.Row = scope.SQLDB().QueryRow(scope.SQL, scope.SQLVars...)
+		} else if rowsResult, ok := result.(*RowsQueryResult); ok {
+			rowsResult.Rows, rowsResult.Error = scope.SQLDB().Query(scope.SQL, scope.SQLVars...)
+		}
+	}
+}
diff --git a/scope.go b/scope.go
index 0a3d6e6f..c36dbb89 100644
--- a/scope.go
+++ b/scope.go
@@ -886,16 +886,22 @@ func (scope *Scope) updatedAttrsWithValues(value interface{}) (results map[strin
 
 func (scope *Scope) row() *sql.Row {
 	defer scope.trace(NowFunc())
+
+	result := &RowQueryResult{}
+	scope.InstanceSet("row_query_result", result)
 	scope.callCallbacks(scope.db.parent.callbacks.rowQueries)
-	scope.prepareQuerySQL()
-	return scope.SQLDB().QueryRow(scope.SQL, scope.SQLVars...)
+
+	return result.Row
 }
 
 func (scope *Scope) rows() (*sql.Rows, error) {
 	defer scope.trace(NowFunc())
+
+	result := &RowsQueryResult{}
+	scope.InstanceSet("row_query_result", result)
 	scope.callCallbacks(scope.db.parent.callbacks.rowQueries)
-	scope.prepareQuerySQL()
-	return scope.SQLDB().Query(scope.SQL, scope.SQLVars...)
+
+	return result.Rows, result.Error
 }
 
 func (scope *Scope) initialize() *Scope {