Function rewriteSelectLimit

backend/plugin/db/pg/query.go:251–278

rewriteSelectLimit adds or adjusts the LIMIT clause of a SELECT statement using byte-offset positions from the omni AST to surgically edit the original SQL text.

(sql string, sel *ast.SelectStmt, limitCount int)

// rewriteSelectLimit adds or adjusts the LIMIT clause of a SELECT statement
// using byte-offset positions from the omni AST to surgically edit the original SQL text.
func rewriteSelectLimit(sql string, sel *ast.SelectStmt, limitCount int) (string, error) {
	if sel.LimitCount != nil {
		// Already has LIMIT — replace the value if ours is lower.
		existingLimit := extractIntFromNode(sel.LimitCount)
		if existingLimit > 0 && existingLimit <= limitCount {
			return sql, nil // existing limit is already lower or equal, keep it
		}
		loc := nodeLocOf(sel.LimitCount)
		if loc.Start >= 0 && loc.End > loc.Start && loc.End <= len(sql) {
			return sql[:loc.Start] + fmt.Sprintf("%d", limitCount) + sql[loc.End:], nil
		}
		// LimitCount is a non-constant expression (e.g. LIMIT $1, LIMIT (1+2)).
		// Cannot safely rewrite in-place; let the caller fall back to CTE wrapper.
		return "", errors.Errorf("cannot rewrite non-constant LIMIT expression")
	}

	// No LIMIT clause — find the right insertion point.
	// PostgreSQL grammar order: ... ORDER BY ... LIMIT ... FOR UPDATE ...
	// LIMIT goes BEFORE FOR UPDATE but AFTER everything else.
	insertPos, beforeLocking := findLimitInsertPosition(sel)
	if beforeLocking {
		// Inserting at the start of FOR UPDATE/SHARE. The original whitespace
		// before FOR becomes the separator before LIMIT; we add a trailing
		// space to separate the limit value from FOR.
		return sql[:insertPos] + fmt.Sprintf("LIMIT %d ", limitCount) + sql[insertPos:], nil
	}
	return sql[:insertPos] + fmt.Sprintf(" LIMIT %d", limitCount) + sql[insertPos:], nil
}

// findLimitInsertPosition returns the byte offset where " LIMIT N" should be inserted,
// and whether the insertion is before a locking clause (FOR UPDATE/SHARE).