Skip to content

Add query time to logger #3010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Jun 3, 2025
8 changes: 4 additions & 4 deletions server/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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.
Expand Down
77 changes: 77 additions & 0 deletions server/handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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
}
1 change: 1 addition & 0 deletions sql/log.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,5 @@ const (
ConnectionIdLogField = "connectionID"
ConnectionDbLogField = "connectionDb"
ConnectTimeLogKey = "connectTime"
QueryTimeLogKey = "queryTime"
)