Thank you for working on this! I have a quick question. This PR seems like a very explicit fix for #74413. So is the plan to have a more general solution for #72045 eventually?

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @cameronnunez and @knz)

pkg/server/drain_test.go, line 143 at r1 (raw file):

	// Check that in-memory data was flushed into system tables during the drain.
	// Verify that the statement stats are in the reported stats pool.
	stats, err = ts.GetScrubbedReportingStats(ctx)

the reporting stats pool is an in-memory pool, so this means GetScrubbedReportingStats() doesn't actually read from the system table at all. My preference here would be a bit explicit and to ensure that the stats is flushed into system by reading them back. This is because the flush is only done in the best-effort basis and if it the flush fails, we simply just log a warning and it won't be retried. So this means it is possible for the stats to end up in the reporting pool but not in the system table.

Good question. As of now, within SQL Observability, SQL Stats is the only subsystem that require this hook. Though as we build up our observability stack, we will have more similar subsystems in the future that requires this type of hook.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @cameronnunez and @knz)

pkg/server/drain_test.go, line 140 at r2 (raw file):

	// Issue a drain.
	drainCtx.sendDrainNoShutdown()

nit: i would do a sanity check before we drain the node to ensure that we have nothing in the system table.

pkg/server/drain_test.go, line 149 at r2 (raw file):

	require.NotEqualf(t,
		func() int {
			q := sqlDB.Query(t, `SELECT count(*) FROM system.statement_statistics WHERE node_id = 1`)

nit: sqlDB.CheckQueryResults can come in pretty handy here :)

Gotcha yeah I think it makes sense to eventually get to the general solution then.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @Azhng and @knz)

pkg/server/drain_test.go, line 143 at r1 (raw file):

Done.

pkg/server/drain_test.go, line 140 at r2 (raw file):

Done.

pkg/server/drain_test.go, line 149 at r2 (raw file):

Would've gone that route but figured it would be best just to check for the count being non-zero so as to future-proof this test; if we record more stats in the future, this test would fail because the count would increase. What're your thoughts on this?

Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @knz)

pkg/server/drain_test.go, line 149 at r2 (raw file):

Hmm how about SELECT count(*) > 0 FROM ... and then just assert the output is true in the CheckQueryResults ?

simple, elegant, wow, amaze 🐶

Reviewed 2 of 3 files at r3, 1 of 1 files at r4, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (and 1 stale) (waiting on @cameronnunez)