getStatementWithResultLimitInline rewrites the single statement so that its outermost query returns at most limitCount rows. The rewrite splices the replacement into the original statement text via omni AST byte offsets (ast.Loc), preserving the original formatting. The output shape mirrors the lega
(singleStatement string, limitCount int)
| 45 | // fired on them). Any error makes the caller fall back to wrapping the |
| 46 | // statement in SELECT * FROM (...) LIMIT n. |
| 47 | func getStatementWithResultLimitInline(singleStatement string, limitCount int) (string, error) { |
| 48 | trimmed := strings.TrimRightFunc(singleStatement, utils.IsSpaceOrSemicolon) |
| 49 | |
| 50 | file, err := omniparser.Parse(trimmed) |
| 51 | if err != nil { |
| 52 | return "", errors.Wrapf(err, stmtErrFmt, singleStatement) |
| 53 | } |
| 54 | if len(file.Stmts) != 1 { |
| 55 | return "", errors.Errorf("expected exactly 1 statement, got %d", len(file.Stmts)) |
| 56 | } |
| 57 | |
| 58 | sel := findRightMostSelect(file.Stmts[0]) |
| 59 | if sel == nil { |
| 60 | // Not a query statement; return the normalized text unchanged. |
| 61 | return trimmed + "\n;", nil |
| 62 | } |
| 63 | |
| 64 | // LIMIT <n> [OFFSET <m>]: replace n. |
| 65 | if sel.Limit != nil { |
| 66 | spliced, err := spliceLimitNumber(trimmed, omniast.NodeLoc(sel.Limit), limitCount) |
| 67 | if err != nil { |
| 68 | return "", errors.Wrapf(err, stmtErrFmt, singleStatement) |
| 69 | } |
| 70 | return spliced + "\n;", nil |
| 71 | } |
| 72 | |
| 73 | // [OFFSET <m>] FETCH [FIRST|NEXT] <n> ...: replace the fetch count n. |
| 74 | if sel.Fetch != nil { |
| 75 | spliced, err := spliceLimitNumber(trimmed, omniast.NodeLoc(sel.Fetch.Count), limitCount) |
| 76 | if err != nil { |
| 77 | return "", errors.Wrapf(err, stmtErrFmt, singleStatement) |
| 78 | } |
| 79 | return spliced + "\n;", nil |
| 80 | } |
| 81 | |
| 82 | // A bare OFFSET without LIMIT/FETCH cannot take an appended LIMIT clause |
| 83 | // (LIMIT must precede OFFSET); the legacy parser rejected this shape, so |
| 84 | // keep routing it to the caller's fallback. |
| 85 | if sel.Offset != nil { |
| 86 | return "", errors.Errorf("statement has OFFSET without LIMIT/FETCH: %s", singleStatement) |
| 87 | } |
| 88 | |
| 89 | // TOP <n>: replace n. |
| 90 | if sel.Top != nil { |
| 91 | spliced, err := spliceLimitNumber(trimmed, omniast.NodeLoc(sel.Top), limitCount) |
| 92 | if err != nil { |
| 93 | return "", errors.Wrapf(err, stmtErrFmt, singleStatement) |
| 94 | } |
| 95 | return spliced + "\n;", nil |
| 96 | } |
| 97 | |
| 98 | // No limiting clause: append after the end of the select. |
| 99 | end := sel.Loc.End |
| 100 | if !sel.Loc.IsValid() || end > len(trimmed) { |
| 101 | return "", errors.Errorf("invalid select location %v for statement: %s", sel.Loc, singleStatement) |
| 102 | } |
| 103 | return trimmed[:end] + fmt.Sprintf(" LIMIT %d", limitCount) + trimmed[end:] + "\n;", nil |
| 104 | } |
no test coverage detected