diff --git a/agents/drivers/kingbase-go/main.go b/agents/drivers/kingbase-go/main.go index 23d45fb609..f8c5698fd2 100644 --- a/agents/drivers/kingbase-go/main.go +++ b/agents/drivers/kingbase-go/main.go @@ -10,6 +10,7 @@ import ( "fmt" "net/url" "os" + "regexp" "strings" "sync" "time" @@ -590,7 +591,46 @@ func (s *server) executeQuery(opts queryOptions) (queryResult, error) { start := time.Now() sqlText := trimStatementSQL(opts.SQL) if isQuerySQL(sqlText) { - rows, conn, cancel, err := s.queryRows(sqlText, opts.Schema, opts.TimeoutSecs) + maxRows := opts.MaxRows + if maxRows <= 0 { + maxRows = defaultMaxRows + } + // DBX core fuses `LIMIT n [OFFSET m]` (and the user's ORDER BY stays) + // into the SQL for pagination. For KingbaseES/Postgres the fused LIMIT + // makes the planner pick a nested-loop plan "optimized for N rows" whose + // repeated pub_code bitmap scans can take minutes, and the ORDER BY forces + // a blocking sort. The fast path is the one the COUNT(*) wrapper already + // gets: no LIMIT, no ORDER BY -> hash join + HashAggregate. So strip the + // trailing LIMIT/OFFSET, wrap the bare statement in a derived table, and + // re-apply ORDER BY + LIMIT/OFFSET on the outside: the inner subquery + // keeps the fast plan while the server sorts the materialized groups and + // pages correctly (ordered results, working OFFSET). If the rewritten + // query fails, fall back to the bare unordered execution. + execSQL := sqlText + if raw, limit, offset, hadLimit := stripTrailingLimit(sqlText); hadLimit { + if bare, orderClause, hadOrder := stripTrailingOrderBy(raw); hadOrder { + pageSize := limit + if pageSize <= 0 { + pageSize = maxRows + } + orderedSQL := buildOrderedPagedQuery(bare, orderClause, pageSize+1, offset) + rows, conn, cancel, err := s.queryRows(orderedSQL, opts.Schema, opts.TimeoutSecs) + if err == nil { + result, rerr := readRows(rows, maxRows) + _ = rows.Close() + _ = conn.Close() + s.endOperation(cancel) + if rerr == nil { + result.ExecutionTimeMS = time.Since(start).Milliseconds() + return result, nil + } + } + execSQL = bare + } else { + execSQL = raw + } + } + rows, conn, cancel, err := s.queryRows(execSQL, opts.Schema, opts.TimeoutSecs) if err != nil { return queryResult{}, err } @@ -599,10 +639,6 @@ func (s *server) executeQuery(opts queryOptions) (queryResult, error) { _ = conn.Close() s.endOperation(cancel) }() - maxRows := opts.MaxRows - if maxRows <= 0 { - maxRows = defaultMaxRows - } result, err := readRows(rows, maxRows) result.ExecutionTimeMS = time.Since(start).Milliseconds() return result, err @@ -623,6 +659,79 @@ func (s *server) executeQuery(opts queryOptions) (queryResult, error) { return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil } +// trailingLimitRe matches a DBX-appended `LIMIT n [OFFSET m]` suffix at the end +// of a statement (optional trailing semicolon/whitespace). It only matches a +// top-level trailing LIMIT, so a LIMIT inside a subquery is preserved. +var trailingLimitRe = regexp.MustCompile(`(?is)\s+LIMIT\s+(\d+)(?:\s+OFFSET\s+(\d+))?\s*;?\s*$`) + +// stripTrailingLimit removes a trailing `LIMIT n [OFFSET m]` from the statement +// and reports the captured limit/offset (zero when absent). +func stripTrailingLimit(sqlText string) (stripped string, limit, offset int, hadLimit bool) { + loc := trailingLimitRe.FindStringSubmatchIndex(sqlText) + if loc == nil { + return sqlText, 0, 0, false + } + limit = atoiSafe(sqlText[loc[2]:loc[3]]) + if loc[4] >= 0 { // optional OFFSET group matched + offset = atoiSafe(sqlText[loc[4]:loc[5]]) + } + return strings.TrimSpace(sqlText[:loc[0]]), limit, offset, true +} + +func atoiSafe(s string) int { + n := 0 + for _, c := range s { + if c < '0' || c > '9' { + return n + } + n = n*10 + int(c-'0') + } + return n +} + +// stripTrailingOrderBy splits the last top-level `ORDER BY ...` clause off the +// statement by scanning left-to-right and tracking parenthesis depth, so an +// ORDER BY inside a subquery is preserved. It does not parse string/identifier +// quotes, so it is only applied to DBX-paginated statements (which already had +// a trailing LIMIT stripped) where the rewrite is the point. +func stripTrailingOrderBy(sqlText string) (bare, orderClause string, hadOrder bool) { + upper := strings.ToUpper(sqlText) + idx := -1 + depth := 0 + for i := 0; i+9 <= len(upper); i++ { + switch sqlText[i] { + case '(': + depth++ + case ')': + if depth > 0 { + depth-- + } + } + if depth == 0 && upper[i:i+9] == " ORDER BY" { + idx = i + } + } + if idx < 0 { + return sqlText, "", false + } + return strings.TrimSpace(sqlText[:idx]), strings.TrimSpace(sqlText[idx:]), true +} + +// orderByQualifierRe matches a `qualifier.` prefix in an ORDER BY column +// reference, so `t.modelKey` can be rewritten to `modelKey` when the ORDER BY is +// moved outside the derived table (which exposes plain column names). +var orderByQualifierRe = regexp.MustCompile(`(?i)\b[a-z_][a-z0-9_]*\.`) + +// buildOrderedPagedQuery wraps the bare aggregate in a derived table and applies +// ORDER BY + LIMIT/OFFSET on the outside. The inner subquery (no ORDER BY) keeps +// the fast hash-join + HashAggregate plan; the outer sorts the materialized +// groups and pages them, so results stay ordered and OFFSET pagination works. +// limit is expected to be pageSize+1 so truncation can be detected. +func buildOrderedPagedQuery(bare, orderClause string, limit, offset int) string { + order := orderByQualifierRe.ReplaceAllString(orderClause, "") + return fmt.Sprintf("SELECT * FROM (%s) dbx_p %s LIMIT %d OFFSET %d", bare, order, limit, offset) +} + func (s *server) queryRows(sqlText string, schema string, timeoutSecs int) (*sql.Rows, *sql.Conn, context.CancelFunc, error) { conn, ctx, cancel, err := s.operationConn(schema, timeoutSecs) if err != nil {