|
| 1 | +"""SQL formatting helpers for the query editor.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import sqlparse |
| 6 | + |
| 7 | +from sqlit.domains.query.app.multi_statement import split_statements |
| 8 | + |
| 9 | + |
| 10 | +def _format_fragment(sql: str) -> str: |
| 11 | + return sqlparse.format( |
| 12 | + sql, |
| 13 | + reindent=True, |
| 14 | + keyword_case="upper", |
| 15 | + use_space_around_operators=True, |
| 16 | + indent_width=4, |
| 17 | + ).strip() |
| 18 | + |
| 19 | + |
| 20 | +def format_sql(sql: str) -> str: |
| 21 | + """Format SQL without changing Sqlit's implicit statement boundaries.""" |
| 22 | + if not sql.strip(): |
| 23 | + return sql |
| 24 | + |
| 25 | + statements = split_statements(sql) |
| 26 | + semicolon_statements = sqlparse.split(sql) |
| 27 | + if len(statements) > 1 and len(semicolon_statements) == 1: |
| 28 | + # Sqlit also supports two blank lines as an implicit separator. Format |
| 29 | + # each statement independently so sqlparse cannot collapse that boundary. |
| 30 | + return "\n\n\n".join(_format_fragment(statement) for statement in statements) |
| 31 | + return _format_fragment(sql) |
| 32 | + |
| 33 | + |
| 34 | +def _location_to_offset(text: str, location: tuple[int, int]) -> int: |
| 35 | + row, column = location |
| 36 | + lines = text.splitlines(keepends=True) |
| 37 | + if row >= len(lines): |
| 38 | + return len(text) |
| 39 | + return min(sum(len(line) for line in lines[:row]) + column, len(text)) |
| 40 | + |
| 41 | + |
| 42 | +def _offset_to_location(text: str, offset: int) -> tuple[int, int]: |
| 43 | + prefix = text[: max(0, min(offset, len(text)))] |
| 44 | + row = prefix.count("\n") |
| 45 | + last_newline = prefix.rfind("\n") |
| 46 | + column = len(prefix) if last_newline < 0 else len(prefix) - last_newline - 1 |
| 47 | + return row, column |
| 48 | + |
| 49 | + |
| 50 | +def remap_cursor_after_format( |
| 51 | + original: str, |
| 52 | + formatted: str, |
| 53 | + cursor: tuple[int, int], |
| 54 | +) -> tuple[int, int]: |
| 55 | + """Keep the cursor beside the same non-whitespace token after formatting.""" |
| 56 | + original_offset = _location_to_offset(original, cursor) |
| 57 | + token_count = sum(not char.isspace() for char in original[:original_offset]) |
| 58 | + if token_count == 0: |
| 59 | + return 0, 0 |
| 60 | + |
| 61 | + seen = 0 |
| 62 | + formatted_offset = len(formatted) |
| 63 | + for index, char in enumerate(formatted): |
| 64 | + if not char.isspace(): |
| 65 | + seen += 1 |
| 66 | + if seen == token_count: |
| 67 | + formatted_offset = index + 1 |
| 68 | + break |
| 69 | + return _offset_to_location(formatted, formatted_offset) |
0 commit comments