stmtdiagnostics: use background context when building the bundle

When the context is canceled, we still want to build the bundle as best
as possible. Over in 532274b we
introduced the usage of the background context in order to insert the
bundle into the system tables, but we still built the bundle with the
canceled context. This commit fixes that oversight - in particular, we
should now get `env.sql` correctly.

Epic: None

Release note: None

Author: yuzefovich

pkg/sql/instrumentation.go

@@ -424,8 +424,27 @@ func (ih *instrumentationHelper) Finish(
 			if pwe, ok := retPayload.(payloadWithError); ok {
 				payloadErr = pwe.errorCause()
 			}
+			bundleCtx := ctx
+			if bundleCtx.Err() != nil {
+				// The only two possible errors on the context are the context
+				// cancellation or the context deadline being exceeded. The
+				// former seems more likely, and the cancellation is most likely
+				// to have occurred due to a statement timeout, so we still want
+				// to proceed with saving the statement bundle. Thus, we
+				// override the canceled context, but first we'll log the error
+				// as a warning.
+				log.Warningf(
+					bundleCtx, "context has an error when saving the bundle, proceeding "+
+						"with the background one (with deadline of 10 seconds): %v", bundleCtx.Err(),
+				)
+				// We want to be conservative, so we add a deadline of 10
+				// seconds on top of the background context.
+				var cancel context.CancelFunc
+				bundleCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second) // nolint:context
+				defer cancel()
+			}
 			bundle = buildStatementBundle(
-				ctx, ih.explainFlags, cfg.DB, ie.(*InternalExecutor), stmtRawSQL, &p.curPlan,
+				bundleCtx, ih.explainFlags, cfg.DB, ie.(*InternalExecutor), stmtRawSQL, &p.curPlan,
 				ob.BuildString(), trace, placeholders, res.Err(), payloadErr, retErr,
 				&p.extendedEvalCtx.Settings.SV,
 			)
@@ -434,7 +453,7 @@ func (ih *instrumentationHelper) Finish(
 			// to the current user and aren't included into the bundle.
 			warnings = append(warnings, bundle.errorStrings...)
 			bundle.insert(
-				ctx, ih.fingerprint, ast, cfg.StmtDiagnosticsRecorder, ih.diagRequestID, ih.diagRequest,
+				bundleCtx, ih.fingerprint, ast, cfg.StmtDiagnosticsRecorder, ih.diagRequestID, ih.diagRequest,
 			)
 			telemetry.Inc(sqltelemetry.StatementDiagnosticsCollectedCounter)
 		}

pkg/sql/stmtdiagnostics/statement_diagnostics.go

@@ -496,8 +496,6 @@ func (r *Registry) ShouldCollectDiagnostics(
 
 // InsertStatementDiagnostics inserts a trace into system.statement_diagnostics.
 //
-// traceJSON is either DNull (when collectionErr should not be nil) or a *DJSON.
-//
 // If requestID is not zero, it also marks the request as completed in
 // system.statement_diagnostics_requests. If requestID is zero, a new entry is
 // inserted.
@@ -514,24 +512,6 @@ func (r *Registry) InsertStatementDiagnostics(
 	collectionErr error,
 ) (CollectedInstanceID, error) {
 	var diagID CollectedInstanceID
-	if ctx.Err() != nil {
-		// The only two possible errors on the context are the context
-		// cancellation or the context deadline being exceeded. The former seems
-		// more likely, and the cancellation is most likely to have occurred due
-		// to a statement timeout, so we still want to proceed with saving the
-		// statement bundle. Thus, we override the canceled context, but first
-		// we'll log the error as a warning.
-		log.Warningf(
-			ctx, "context has an error when saving the bundle, proceeding "+
-				"with the background one (with deadline of 10 seconds): %v", ctx.Err(),
-		)
-		// We want to be conservative, so we add a deadline of 10 seconds on top
-		// of the background context.
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second) // nolint:context
-		defer cancel()
-	}
-
 	err := r.db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
 		txn.KV().SetDebugName("stmt-diag-insert-bundle")
 		if requestID != 0 {