From 4671b36568561cd5683d6a648fe8cec282ddff68 Mon Sep 17 00:00:00 2001
From: qiangxue <qiang.xue@gmail.com>
Date: Tue, 10 Oct 2017 21:57:04 -0400
Subject: [PATCH] added db.PerfFunc

---
 README.md |  5 ++++-
 db.go     |  9 +++++++++
 query.go  | 30 +++++++++++++++++++-----------
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index f798377..e2489ac 100644
--- a/README.md
+++ b/README.md
@@ -673,7 +673,10 @@ func main() {
 )
 ```
 
-And the following example shows how to use the `ozzo-log` package which allows logging message severities and categories
+You can also configure `DB.PerfFunc` to capture the SQL statement execution times. Each time when a SQL statement
+is executed or queried, this function will be called with the time used. This allows you to profile your DB performance.
+
+The following example shows how to use the `ozzo-log` package which allows logging message severities and categories
 and sending logged messages to different targets (e.g. files, console window, network).
 
 ```go
diff --git a/db.go b/db.go
index bba7b1e..2ab8cf7 100644
--- a/db.go
+++ b/db.go
@@ -19,6 +19,13 @@ type (
 	// are provided, they will be passed to fmt.Sprintf() to generate the log message.
 	LogFunc func(format string, a ...interface{})
 
+	// PerfFunc is called when a query finishes execution.
+	// The query execution time is passed to this function so that the DB performance
+	// can be profiled. The "ns" parameter gives the number of nanoseconds that the
+	// SQL statement takes to execute, while the "execute" parameter indicates whether
+	// the SQL statement is executed or queried (usually SELECT statements).
+	PerfFunc func(ns int64, sql string, execute bool)
+
 	// BuilderFunc creates a Builder instance using the given DB instance and Executor.
 	BuilderFunc func(*DB, Executor) Builder
 
@@ -31,6 +38,8 @@ type (
 		FieldMapper FieldMapFunc
 		// LogFunc logs the SQL statements being executed. Defaults to nil, meaning no logging.
 		LogFunc LogFunc
+		// PerfFunc logs the SQL execution time. Defaults to nil, meaning no performance profiling.
+		PerfFunc PerfFunc
 
 		sqlDB      *sql.DB
 		driverName string
diff --git a/query.go b/query.go
index f3f8ffb..df19e62 100644
--- a/query.go
+++ b/query.go
@@ -44,6 +44,8 @@ type Query struct {
 	LastError error
 	// LogFunc is used to log the SQL statement being executed.
 	LogFunc LogFunc
+	// PerfFunc is used to log the SQL execution time. It is ignored if nil.
+	PerfFunc PerfFunc
 }
 
 // NewQuery creates a new Query with the given SQL statement.
@@ -57,6 +59,7 @@ func NewQuery(db *DB, executor Executor, sql string) *Query {
 		params:       Params{},
 		FieldMapper:  db.FieldMapper,
 		LogFunc:      db.LogFunc,
+		PerfFunc:     db.PerfFunc,
 	}
 }
 
@@ -90,11 +93,20 @@ func (q *Query) logSQL() string {
 
 // log logs a message for the currently executed SQL statement.
 func (q *Query) log(start time.Time, execute bool) {
-	t := float64(time.Now().Sub(start).Nanoseconds()) / 1e6
-	if execute {
-		q.LogFunc("[%.2fms] Execute SQL: %v", t, q.logSQL())
-	} else {
-		q.LogFunc("[%.2fms] Query SQL: %v", t, q.logSQL())
+	if q.LogFunc == nil && q.PerfFunc == nil {
+		return
+	}
+	ns := time.Now().Sub(start).Nanoseconds()
+	s := q.logSQL()
+	if q.LogFunc != nil {
+		if execute {
+			q.LogFunc("[%.2fms] Execute SQL: %v", float64(ns)/1e6, s)
+		} else {
+			q.LogFunc("[%.2fms] Query SQL: %v", float64(ns)/1e6, s)
+		}
+	}
+	if q.PerfFunc != nil {
+		q.PerfFunc(ns, s, execute)
 	}
 }
 
@@ -154,9 +166,7 @@ func (q *Query) Execute() (result sql.Result, err error) {
 		return
 	}
 
-	if q.LogFunc != nil {
-		defer q.log(time.Now(), true)
-	}
+	defer q.log(time.Now(), true)
 
 	if q.stmt == nil {
 		result, err = q.executor.Exec(q.rawSQL, params...)
@@ -224,9 +234,7 @@ func (q *Query) Rows() (rows *Rows, err error) {
 		return
 	}
 
-	if q.LogFunc != nil {
-		defer q.log(time.Now(), false)
-	}
+	defer q.log(time.Now(), false)
 
 	var rr *sql.Rows
 	if q.stmt == nil {