| 224 | } |
| 225 | |
| 226 | func getStatementWithResultLimitInline(statement string, limitCount int) (string, error) { |
| 227 | if strings.TrimSpace(statement) == "" { |
| 228 | return "", errors.New("empty statement") |
| 229 | } |
| 230 | |
| 231 | stmts, err := pgparser.ParsePg(statement) |
| 232 | if err != nil { |
| 233 | return "", errors.Wrap(err, "failed to parse statement") |
| 234 | } |
| 235 | |
| 236 | if len(stmts) != 1 { |
| 237 | return "", errors.Errorf("expected exactly one statement, got %d", len(stmts)) |
| 238 | } |
| 239 | |
| 240 | sel, ok := stmts[0].AST.(*ast.SelectStmt) |
| 241 | if !ok { |
| 242 | // Non-SELECT statement, return as-is. |
| 243 | return statement, nil |
| 244 | } |
| 245 | |
| 246 | return rewriteSelectLimit(statement, sel, limitCount) |
| 247 | } |
| 248 | |
| 249 | // rewriteSelectLimit adds or adjusts the LIMIT clause of a SELECT statement |
| 250 | // using byte-offset positions from the omni AST to surgically edit the original SQL text. |