From 4e559eee5d1e8447eb22ff896510112ec8bf17b7 Mon Sep 17 00:00:00 2001
From: Lunny Xiao <xiaolunwen@gmail.com>
Date: Sat, 29 Feb 2020 08:59:03 +0800
Subject: [PATCH] Fix bug

---
 core/db.go | 56 +++++++++++++++++++++++++++++++-----------------------
 core/tx.go | 49 +++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 26 deletions(-)

diff --git a/core/db.go b/core/db.go
index 1e25d007..ebd3f689 100644
--- a/core/db.go
+++ b/core/db.go
@@ -124,19 +124,23 @@ func (db *DB) reflectNew(typ reflect.Type) reflect.Value {
 // QueryContext overwrites sql.DB.QueryContext
 func (db *DB) QueryContext(ctx context.Context, query string, args ...interface{}) (*Rows, error) {
 	start := time.Now()
-	db.Logger.BeforeSQL(log.LogContext{
-		Ctx:  ctx,
-		SQL:  query,
-		Args: args,
-	})
+	if db.Logger != nil {
+		db.Logger.BeforeSQL(log.LogContext{
+			Ctx:  ctx,
+			SQL:  query,
+			Args: args,
+		})
+	}
 	rows, err := db.DB.QueryContext(ctx, query, args...)
-	db.Logger.AfterSQL(log.LogContext{
-		Ctx:         ctx,
-		SQL:         query,
-		Args:        args,
-		ExecuteTime: time.Now().Sub(start),
-		Err:         err,
-	})
+	if db.Logger != nil {
+		db.Logger.AfterSQL(log.LogContext{
+			Ctx:         ctx,
+			SQL:         query,
+			Args:        args,
+			ExecuteTime: time.Now().Sub(start),
+			Err:         err,
+		})
+	}
 	if err != nil {
 		if rows != nil {
 			rows.Close()
@@ -242,19 +246,23 @@ func (db *DB) ExecStructContext(ctx context.Context, query string, st interface{
 
 func (db *DB) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
 	start := time.Now()
-	db.Logger.BeforeSQL(log.LogContext{
-		Ctx:  ctx,
-		SQL:  query,
-		Args: args,
-	})
+	if db.Logger != nil {
+		db.Logger.BeforeSQL(log.LogContext{
+			Ctx:  ctx,
+			SQL:  query,
+			Args: args,
+		})
+	}
 	res, err := db.DB.ExecContext(ctx, query, args...)
-	db.Logger.AfterSQL(log.LogContext{
-		Ctx:         ctx,
-		SQL:         query,
-		Args:        args,
-		ExecuteTime: time.Now().Sub(start),
-		Err:         err,
-	})
+	if db.Logger != nil {
+		db.Logger.AfterSQL(log.LogContext{
+			Ctx:         ctx,
+			SQL:         query,
+			Args:        args,
+			ExecuteTime: time.Now().Sub(start),
+			Err:         err,
+		})
+	}
 	return res, err
 }
 
diff --git a/core/tx.go b/core/tx.go
index a56b7006..614213c9 100644
--- a/core/tx.go
+++ b/core/tx.go
@@ -7,6 +7,9 @@ package core
 import (
 	"context"
 	"database/sql"
+	"time"
+
+	"xorm.io/xorm/log"
 )
 
 type Tx struct {
@@ -64,7 +67,7 @@ func (tx *Tx) ExecMapContext(ctx context.Context, query string, mp interface{})
 	if err != nil {
 		return nil, err
 	}
-	return tx.Tx.ExecContext(ctx, query, args...)
+	return tx.ExecContext(ctx, query, args...)
 }
 
 func (tx *Tx) ExecMap(query string, mp interface{}) (sql.Result, error) {
@@ -76,7 +79,29 @@ func (tx *Tx) ExecStructContext(ctx context.Context, query string, st interface{
 	if err != nil {
 		return nil, err
 	}
-	return tx.Tx.ExecContext(ctx, query, args...)
+	return tx.ExecContext(ctx, query, args...)
+}
+
+func (tx *Tx) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	start := time.Now()
+	if tx.db.Logger != nil {
+		tx.db.Logger.BeforeSQL(log.LogContext{
+			Ctx:  ctx,
+			SQL:  query,
+			Args: args,
+		})
+	}
+	res, err := tx.Tx.ExecContext(ctx, query, args...)
+	if tx.db.Logger != nil {
+		tx.db.Logger.AfterSQL(log.LogContext{
+			Ctx:         ctx,
+			SQL:         query,
+			Args:        args,
+			ExecuteTime: time.Now().Sub(start),
+			Err:         err,
+		})
+	}
+	return res, err
 }
 
 func (tx *Tx) ExecStruct(query string, st interface{}) (sql.Result, error) {
@@ -84,8 +109,28 @@ func (tx *Tx) ExecStruct(query string, st interface{}) (sql.Result, error) {
 }
 
 func (tx *Tx) QueryContext(ctx context.Context, query string, args ...interface{}) (*Rows, error) {
+	start := time.Now()
+	if tx.db.Logger != nil {
+		tx.db.Logger.BeforeSQL(log.LogContext{
+			Ctx:  ctx,
+			SQL:  query,
+			Args: args,
+		})
+	}
 	rows, err := tx.Tx.QueryContext(ctx, query, args...)
+	if tx.db.Logger != nil {
+		tx.db.Logger.AfterSQL(log.LogContext{
+			Ctx:         ctx,
+			SQL:         query,
+			Args:        args,
+			ExecuteTime: time.Now().Sub(start),
+			Err:         err,
+		})
+	}
 	if err != nil {
+		if rows != nil {
+			rows.Close()
+		}
 		return nil, err
 	}
 	return &Rows{rows, tx.db}, nil