session: drain result sets in --initialize-sql-file before closing

[ljluestc]

What problem does this PR solve?

Issue Number: close #63450

Problem Summary:

When bootstrapping TiDB with --initialize-sql-file, SELECT statements that call functions with side effects do not work. For example:

SELECT audit_log_create_filter('all', '{}');
SELECT audit_log_create_rule('test@%', 'all');

After bootstrap, SELECT * FROM mysql.audit_log_filters returns an empty result. The workaround is to use DO instead of SELECT, but SELECT should also work.

Root cause: In doBootstrapSQLFile (pkg/session/bootstrap.go), result sets from SELECT statements were closed via rs.Close() without being drained. TiDB evaluates expressions lazily during row iteration (rs.Next()). When a SELECT calls a function with side effects, the function body is only executed when rows are fetched. Since Close() was called without any Next() calls, the function was never evaluated and its side effects never occurred.

What changed and how does it work?

pkg/session/bootstrap.go — In doBootstrapSQLFile, replace the immediate rs.Close() with a drain loop that calls rs.Next() until all rows are consumed before closing. This ensures all row-level expression evaluation (including side-effect functions) completes.

Before (broken):

if rs != nil {
    // I don't believe we need to drain the result-set in bootstrap mode
    // but if required we can do this here in future.
    if err := rs.Close(); err != nil {
        logutil.BgLogger().Fatal("unable to close result", zap.Error(err))
    }
}

After (fixed):

if rs != nil {
    // Drain the result set before closing. For SELECT statements that
    // call functions with side effects (e.g. audit_log_create_filter),
    // the function is evaluated lazily during row iteration. Without
    // draining, the side effects never take place. See #63450.
    chk := rs.NewChunk(nil)
    for {
        if err := rs.Next(ctx, chk); err != nil {
            logutil.BgLogger().Warn("InitializeSQLFile result set drain error", zap.Error(err))
            break
        }
        if chk.NumRows() == 0 {
            break
        }
    }
    if err := rs.Close(); err != nil {
        logutil.BgLogger().Fatal("unable to close result", zap.Error(err))
    }
}

pkg/session/test/session_test.go — Add TestInitSQLFileSelectDrain regression test that bootstraps with an init SQL file containing mixed DDL, SELECT, and DML statements, verifying all statements execute correctly after the SELECT result sets are drained.

Check List

Tests

• Unit test
• Integration test
• Manual test (add detailed scripts or steps below)
• No need to test

  • I checked and no code files have been changed.

Unit test (all PASS):

pushd pkg/session/test
go test -run "TestInitSQLFileSelectDrain|TestInitSystemVariable|TestInitUsers|TestEmptyInitSQLFile|TestErrorHappenWhileInit" -count=1 -timeout 120s --tags=intest -v
popd

# Results:
# --- PASS: TestInitSQLFileSelectDrain (0.67s)   <-- new regression test
# --- PASS: TestInitSystemVariable (0.59s)
# --- PASS: TestInitUsers (0.72s)
# --- PASS: TestEmptyInitSQLFile (0.73s)
# --- PASS: TestErrorHappenWhileInit (1.33s)

No integration tests exist for --initialize-sql-file (checked tests/integrationtest/ and tests/ — no matches).

Manual test (enterprise build):

# init.sql
SELECT audit_log_create_filter('all', '{}');
SELECT audit_log_create_rule('test@%', 'all');

# Bootstrap
./tidb-server --initialize-sql-file=./init.sql

# Verify
mysql> SELECT * FROM mysql.audit_log_filters;
# Should return the created filter (previously returned empty)

Side effects

• Performance regression: Consumes more CPU
• Performance regression: Consumes more Memory
• Breaking backward compatibility

Documentation

• Affects user behaviors
• Contains syntax changes
• Contains variable changes
• Contains experimental features
• Changes MySQL compatibility

Release note

Please refer to Release Notes Language Style Guide to write a quality release note.

Fix an issue where `SELECT` statements calling functions with side effects in `--initialize-sql-file` did not execute the functions because the result set was not drained before closing (#63450).

[ti-chi-bot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign benjamin2037 for approval. For more information see the Code Review Process.
Please ensure that each of them provides their approval before proceeding.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS
• pkg/session/OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[ti-chi-bot]

Hi @ljluestc. Thanks for your PR.

I'm waiting for a pingcap member to verify that this patch is reasonable to test. If it is, they should reply with /ok-to-test on its own line. Until that is done, I will not automatically test new commits in this PR, but the usual testing commands by org members will still work. Regular contributors should join the org to skip this step.

Once the patch is verified, the new status will be reflected by the ok-to-test label.

I understand the commands that are listed here.

Details

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.

[pingcap-cla-assistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[tiprow]

Hi @ljluestc. Thanks for your PR.

PRs from untrusted users cannot be marked as trusted with /ok-to-test in this repo meaning untrusted PR authors can never trigger tests themselves. Collaborators can still trigger tests on the PR using /test all.

I understand the commands that are listed here.

Details

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.