Actually emit the correct query time, not just set up the field

Author: timsehn

server/handler.go

@@ -447,29 +447,24 @@ func (h *Handler) doQuery(
 		}
 	}
 
-	// Add query time to logger for this query (after connect time)
-	logger := sqlCtx.GetLogger()
 	if queryStr != "" {
-		logger = logger.WithField(sql.QueryTimeLogKey, start).WithField("query", queryStr)
-	} else {
-		logger = logger.WithField(sql.QueryTimeLogKey, start)
+		sqlCtx.SetLogger(sqlCtx.GetLogger().WithField("query", queryStr))
 	}
-	sqlCtx.SetLogger(logger)
-	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)
 		}
@@ -516,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

@@ -2037,23 +2037,23 @@ func TestHandlerDoQueryIntegration(t *testing.T) {
 	// but won't persist in the session logger (which is the expected behavior)
 }
 
-// TestQueryTimeLoggerBehavior tests that query time is added to logger during query execution
+// TestQueryTimeLoggerBehavior tests that query time field can be added to logger entries
 func TestQueryTimeLoggerBehavior(t *testing.T) {
 	// Test demonstrates that the QueryTimeLogKey constant is available and can be used
-	// for setting up logger fields (which happens in the doQuery method)
+	// for setting up logger fields with time.Now() calls during query execution
 	require.NotEmpty(t, sql.QueryTimeLogKey, "QueryTimeLogKey should not be empty")
 
-	// Create a time value like what would be used in doQuery
+	// Create a time value like what would be used in query logging
 	queryTime := time.Now()
 
-	// Simulate what happens in doQuery: create a logger with query time field
+	// Simulate what happens during query execution: create a logger with query time field
 	baseLogger := &logrus.Entry{
 		Data: make(logrus.Fields),
 	}
 	baseLogger.Data[sql.ConnectTimeLogKey] = time.Now().Add(-1 * time.Minute) // connection established 1 minute ago
 	baseLogger.Data[sql.ConnectionIdLogField] = uint32(123)
 
-	// Add query time field (this is what our modified doQuery does)
+	// Add query time field using time.Now() (this is what our modified handler does)
 	loggerWithQueryTime := baseLogger.WithField(sql.QueryTimeLogKey, queryTime)
 
 	// Verify the query time field was added