Merge pull request #3010 from dolthub/tim/add-query-time-to-logger

Add query time to logger

Author: timsehn

server/handler.go

@@ -450,21 +450,21 @@ func (h *Handler) doQuery(
 	if queryStr != "" {
 		sqlCtx.SetLogger(sqlCtx.GetLogger().WithField("query", queryStr))
 	}
-	sqlCtx.GetLogger().Debugf("Starting query")
+	sqlCtx.GetLogger().WithField(sql.QueryTimeLogKey, time.Now()).Debugf("Starting query")
 
 	finish := observeQuery(sqlCtx, query)
 	defer func() {
 		finish(err)
 	}()
 
-	sqlCtx.GetLogger().Tracef("beginning execution")
+	sqlCtx.GetLogger().WithField(sql.QueryTimeLogKey, time.Now()).Tracef("beginning execution")
 
 	var schema sql.Schema
 	var rowIter sql.RowIter
 	qFlags.Set(sql.QFlagDeferProjections)
 	schema, rowIter, qFlags, err = queryExec(sqlCtx, query, parsed, analyzedPlan, bindings, qFlags)
 	if err != nil {
-		sqlCtx.GetLogger().WithError(err).Warn("error running query")
+		sqlCtx.GetLogger().WithField(sql.QueryTimeLogKey, time.Now()).WithError(err).Warn("error running query")
 		if verboseErrorLogging {
 			fmt.Printf("Err: %+v", err)
 		}
@@ -511,7 +511,7 @@ func (h *Handler) doQuery(
 		sqlCtx.GetLogger().Tracef("returning result %v", r)
 	}
 
-	sqlCtx.GetLogger().Debugf("Query finished in %d ms", time.Since(start).Milliseconds())
+	sqlCtx.GetLogger().WithField(sql.QueryTimeLogKey, time.Now()).Debugf("Query finished in %d ms", time.Since(start).Milliseconds())
 
 	// processedAtLeastOneBatch means we already called callback() at least
 	// once, so no need to call it if RowsAffected == 0.

server/handler_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/dolthub/vitess/go/race"
 	"github.com/dolthub/vitess/go/sqltypes"
 	"github.com/dolthub/vitess/go/vt/proto/query"
+	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -1916,3 +1917,79 @@ func TestStatusVariableComUpdate(t *testing.T) {
 	checkSessionStatVar(t, sess1, "Com_update", uint64(5))
 	checkSessionStatVar(t, sess2, "Com_update", uint64(3))
 }
+
+// TestLoggerFieldsSetup tests that handler properly sets up logger fields including query time
+func TestLoggerFieldsSetup(t *testing.T) {
+	e, pro := setupMemDB(require.New(t))
+	dbFunc := pro.Database
+
+	handler := &Handler{
+		e: e,
+		sm: NewSessionManager(
+			sql.NewContext,
+			testSessionBuilder(pro),
+			sql.NoopTracer,
+			dbFunc,
+			sql.NewMemoryManager(nil),
+			sqle.NewProcessList(),
+			"foo",
+		),
+		readTimeout: time.Second,
+	}
+
+	conn := newConn(1)
+	handler.NewConnection(conn)
+	err := handler.ComInitDB(conn, "test")
+	require.NoError(t, err)
+
+	// Execute a query and verify basic logging setup
+	err = handler.ComQuery(context.Background(), conn, "SELECT 1", dummyCb)
+	require.NoError(t, err)
+
+	// Verify that the session's logger has the expected fields
+	session := handler.sm.session(conn)
+	logger := session.GetLogger()
+	require.NotNil(t, logger, "Session should have a logger")
+
+	// Verify that the logger has the expected fields
+	require.Contains(t, logger.Data, sql.ConnectTimeLogKey, "Logger should contain connect time")
+	require.Contains(t, logger.Data, sql.ConnectionIdLogField, "Logger should contain connection ID")
+
+	// Verify that queryTime is actually used in logs by capturing a log entry
+	var capturedFields logrus.Fields
+	hook := &testHook{fields: &capturedFields}
+	logrus.AddHook(hook)
+	defer logrus.StandardLogger().ReplaceHooks(make(logrus.LevelHooks))
+
+	// Execute a query that will trigger error logging (which includes queryTime)
+	err = handler.ComQuery(context.Background(), conn, "SELECT * FROM nonexistent_table", dummyCb)
+	require.Error(t, err) // This should cause an error log with queryTime
+
+	// Verify that the log entry contained queryTime
+	require.Contains(t, capturedFields, sql.QueryTimeLogKey, "Log entry should contain queryTime field")
+
+	// Verify the values are of correct types
+	connectTime, ok := logger.Data[sql.ConnectTimeLogKey].(time.Time)
+	require.True(t, ok, "Connect time should be a time.Time")
+	require.False(t, connectTime.IsZero(), "Connect time should not be zero")
+
+	connID, ok := logger.Data[sql.ConnectionIdLogField].(uint32)
+	require.True(t, ok, "Connection ID should be a uint32")
+	require.Equal(t, conn.ConnectionID, connID, "Connection ID should match")
+}
+
+// Simple hook to capture log fields for testing
+type testHook struct {
+	fields *logrus.Fields
+}
+
+func (h *testHook) Levels() []logrus.Level {
+	return []logrus.Level{logrus.WarnLevel} // Only capture warning level (error logs)
+}
+
+func (h *testHook) Fire(entry *logrus.Entry) error {
+	if entry.Message == "error running query" {
+		*h.fields = entry.Data
+	}
+	return nil
+}

sql/log.go

@@ -18,4 +18,5 @@ const (
 	ConnectionIdLogField = "connectionID"
 	ConnectionDbLogField = "connectionDb"
 	ConnectTimeLogKey    = "connectTime"
+	QueryTimeLogKey      = "queryTime"
 )