2 files changed, 88 insertions, 0 deletions
diff --git a/internal/dbutil/writer.go b/internal/dbutil/writer.go
new file mode 100644
index 00000000..6b83d537
--- /dev/null
+++ b/internal/dbutil/writer.go
@@ -0,0 +1,35 @@
+// Copyright 2020 The Gogs Authors. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package dbutil
+
+import (
+	"fmt"
+	"io"
+)
+
+// Writer is a wrapper of io.Writer for the gorm.logger.
+type Writer struct {
+	io.Writer
+}
+
+func (w *Writer) Print(v ...interface{}) {
+	if len(v) == 0 {
+		return
+	}
+
+	if len(v) == 1 {
+		fmt.Fprint(w.Writer, v[0])
+		return
+	}
+
+	switch v[0] {
+	case "sql":
+		fmt.Fprintf(w.Writer, "[sql] [%s] [%s] %s %v (%d rows affected)", v[1:]...)
+	case "log":
+		fmt.Fprintf(w.Writer, "[log] [%s] %s", v[1:]...)
+	default:
+		fmt.Fprint(w.Writer, v...)
+	}
+}
diff --git a/internal/dbutil/writer_test.go b/internal/dbutil/writer_test.go
new file mode 100644
index 00000000..33961aa1
--- /dev/null
+++ b/internal/dbutil/writer_test.go
@@ -0,0 +1,53 @@
+// Copyright 2020 The Gogs Authors. All rights reserved.
+// Use of this source code is governed by a MIT-style
+// license that can be found in the LICENSE file.
+
+package dbutil
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWriter_Print(t *testing.T) {
+	tests := []struct {
+		name      string
+		vs        []interface{}
+		expOutput string
+	}{
+		{
+			name: "no values",
+		},
+		{
+			name:      "only one value",
+			vs:        []interface{}{"test"},
+			expOutput: "test",
+		},
+		{
+			name:      "two values",
+			vs:        []interface{}{"test", "output"},
+			expOutput: "testoutput",
+		},
+
+		{
+			name:      "sql",
+			vs:        []interface{}{"sql", "writer.go:65", "1ms", "SELECT * FROM users WHERE user_id = $1", []int{1}, 1},
+			expOutput: "[sql] [writer.go:65] [1ms] SELECT * FROM users WHERE user_id = $1 [1] (1 rows affected)",
+		},
+		{
+			name:      "log",
+			vs:        []interface{}{"log", "writer.go:65", "something"},
+			expOutput: "[log] [writer.go:65] something",
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var buf bytes.Buffer
+			w := &Writer{Writer: &buf}
+			w.Print(test.vs...)
+			assert.Equal(t, test.expOutput, buf.String())
+		})
+	}
+}