Add query formatting from the leader menu

config/keymap.template.json

@@ -245,6 +245,7 @@
         "cancel_operation": "z",
         "change_theme": "t",
         "edit_query_in_editor": "o",
+        "format_query": "p",
         "show_help": "h",
         "telescope": "space",
         "telescope_filter": "/"

sqlit/core/keymap.py

@@ -203,6 +203,13 @@ def _build_leader_commands(self) -> list[LeaderCommandDef]:
                 "Open in editor",
                 "Actions",
             ),
+            LeaderCommandDef(
+                "p",
+                "format_query",
+                "Format Query",
+                "Actions",
+                guard="query_focused",
+            ),
             LeaderCommandDef("h", "show_help", "Help", "Actions"),
             LeaderCommandDef("k", "show_keybinding_editor", "Edit Keybindings", "Actions"),
             LeaderCommandDef("space", "telescope", "Telescope", "Actions"),

sqlit/domains/query/editing/formatting.py

@@ -0,0 +1,69 @@
+"""SQL formatting helpers for the query editor."""
+
+from __future__ import annotations
+
+import sqlparse
+
+from sqlit.domains.query.app.multi_statement import split_statements
+
+
+def _format_fragment(sql: str) -> str:
+    return sqlparse.format(
+        sql,
+        reindent=True,
+        keyword_case="upper",
+        use_space_around_operators=True,
+        indent_width=4,
+    ).strip()
+
+
+def format_sql(sql: str) -> str:
+    """Format SQL without changing Sqlit's implicit statement boundaries."""
+    if not sql.strip():
+        return sql
+
+    statements = split_statements(sql)
+    semicolon_statements = sqlparse.split(sql)
+    if len(statements) > 1 and len(semicolon_statements) == 1:
+        # Sqlit also supports two blank lines as an implicit separator. Format
+        # each statement independently so sqlparse cannot collapse that boundary.
+        return "\n\n\n".join(_format_fragment(statement) for statement in statements)
+    return _format_fragment(sql)
+
+
+def _location_to_offset(text: str, location: tuple[int, int]) -> int:
+    row, column = location
+    lines = text.splitlines(keepends=True)
+    if row >= len(lines):
+        return len(text)
+    return min(sum(len(line) for line in lines[:row]) + column, len(text))
+
+
+def _offset_to_location(text: str, offset: int) -> tuple[int, int]:
+    prefix = text[: max(0, min(offset, len(text)))]
+    row = prefix.count("\n")
+    last_newline = prefix.rfind("\n")
+    column = len(prefix) if last_newline < 0 else len(prefix) - last_newline - 1
+    return row, column
+
+
+def remap_cursor_after_format(
+    original: str,
+    formatted: str,
+    cursor: tuple[int, int],
+) -> tuple[int, int]:
+    """Keep the cursor beside the same non-whitespace token after formatting."""
+    original_offset = _location_to_offset(original, cursor)
+    token_count = sum(not char.isspace() for char in original[:original_offset])
+    if token_count == 0:
+        return 0, 0
+
+    seen = 0
+    formatted_offset = len(formatted)
+    for index, char in enumerate(formatted):
+        if not char.isspace():
+            seen += 1
+            if seen == token_count:
+                formatted_offset = index + 1
+                break
+    return _offset_to_location(formatted, formatted_offset)

sqlit/domains/query/state/query_normal.py

@@ -27,6 +27,7 @@ def _setup_actions(self) -> None:
             "edit_query_in_editor",
             help="Open current query in your terminal editor",
         )
+        self.allows("format_query", help="Format the current query")
         # Vim cursor movement
         self.allows("cursor_left", help="Move cursor left")
         self.allows("cursor_right", help="Move cursor right")

sqlit/domains/query/ui/mixins/query.py

@@ -12,6 +12,7 @@
 from .query_editing_comments import QueryEditingCommentsMixin
 from .query_editing_common import QueryEditingCommonMixin
 from .query_editing_cursor import QueryEditingCursorMixin
+from .query_editing_format import QueryEditingFormatMixin
 from .query_editing_operators import QueryEditingOperatorsMixin
 from .query_editing_selection import QueryEditingSelectionMixin
 from .query_editing_undo import QueryEditingUndoMixin
@@ -29,6 +30,7 @@ class QueryMixin(
     QueryEditingSelectionMixin,
     QueryEditingOperatorsMixin,
     QueryEditingClipboardMixin,
+    QueryEditingFormatMixin,
     QueryEditingCommentsMixin,
     QueryEditingCursorMixin,
     QueryExecutionMixin,

sqlit/domains/query/ui/mixins/query_editing_format.py

@@ -0,0 +1,35 @@
+"""Query formatting action."""
+
+from __future__ import annotations
+
+from sqlit.shared.ui.protocols import QueryMixinHost
+
+
+class QueryEditingFormatMixin:
+    """Format the complete query buffer with sqlparse."""
+
+    def action_format_query(self: QueryMixinHost) -> None:
+        from sqlit.domains.query.editing.formatting import (
+            format_sql,
+            remap_cursor_after_format,
+        )
+
+        original = self.query_input.text
+        if not original.strip():
+            self.notify("Nothing to format", severity="warning")
+            return
+
+        formatted = format_sql(original)
+        if formatted == original:
+            self.notify("Query is already formatted")
+            return
+
+        cursor = remap_cursor_after_format(
+            original,
+            formatted,
+            self.query_input.cursor_location,
+        )
+        self._push_undo_state()
+        self.query_input.text = formatted
+        self.query_input.cursor_location = cursor
+        self.notify("Query formatted")

sqlit/domains/shell/state/machine.py

@@ -34,8 +34,8 @@
     QueryFocusedState,
     QueryInsertModeState,
     QueryNormalModeState,
-    QueryVisualModeState,
     QueryVisualLineModeState,
+    QueryVisualModeState,
 )
 from sqlit.domains.results.state import (
     ResultsFilterActiveState,
@@ -50,7 +50,6 @@
 from sqlit.domains.shell.state.modal_active import ModalActiveState
 from sqlit.domains.shell.state.root import RootState
 
-
 STATE_TO_HELP_SECTION: dict[str, str] = {
     "QueryInsertModeState": "query_insert",
     "AutocompleteActiveState": "query_insert",
@@ -238,6 +237,7 @@ def lk(action: str, menu: str, fallback: str) -> str:
         s.binding(f"{g_key}{lk('execute_query_atomic', 'g', 't')}", "Execute as transaction")
         s.binding(k("show_history", "<backspace>"), "Query history")
         s.binding(k("new_query", "N"), "New query (clear)")
+        s.binding(f"{leader_key}{lk('format_query', 'leader', 'p')}", "Format query")
         s.binding(k("undo", "u"), "Undo")
         s.binding(k("redo", "^r"), "Redo")
         sections.append(s)

sqlit/domains/shell/ui/mixins/ui_leader.py

@@ -109,6 +109,9 @@ def action_leader_change_theme(self: UINavigationMixinHost) -> None:
     def action_leader_edit_query_in_editor(self: UINavigationMixinHost) -> None:
         self._execute_leader_command("edit_query_in_editor")
 
+    def action_leader_format_query(self: UINavigationMixinHost) -> None:
+        self._execute_leader_command("format_query")
+
     def action_leader_toggle_process_worker(self: UINavigationMixinHost) -> None:
         self._execute_leader_command("toggle_process_worker")
 

tests/ui/keybindings/test_keymap_provider.py

@@ -26,6 +26,7 @@ def test_default_keymap_has_expected_leader_commands(self):
         assert keymap.leader("show_help") is not None
         assert keymap.leader("toggle_explorer") is not None
         assert keymap.leader("change_theme") is not None
+        assert keymap.leader("format_query") == "p"
 
     def test_default_keymap_has_expected_action_keys(self):
         """Default keymap should have standard action keys."""

tests/ui/keybindings/test_state_machine.py

@@ -58,7 +58,7 @@ def test_footer_shows_cancel_when_executing(self):
         sm = UIStateMachine()
         ctx = make_context(query_executing=True)
 
-        left, right = sm.get_display_bindings(ctx)
+        left, _right = sm.get_display_bindings(ctx)
         actions = [b.action for b in left]
         assert "cancel_operation" in actions
 
@@ -245,3 +245,45 @@ def test_allowed_when_results_focused(self):
             leader_menu="leader",
         )
         assert sm.check_action(ctx, "leader_edit_query_in_editor") is True
+
+
+class TestFormatQueryLeaderCommand:
+    def test_allowed_when_query_focused(self):
+        sm = UIStateMachine()
+        ctx = make_context(
+            focus="query",
+            leader_pending=True,
+            leader_menu="leader",
+        )
+        assert sm.check_action(ctx, "leader_format_query") is True
+
+    def test_blocked_when_explorer_focused(self):
+        sm = UIStateMachine()
+        ctx = make_context(
+            focus="explorer",
+            leader_pending=True,
+            leader_menu="leader",
+        )
+        assert sm.check_action(ctx, "leader_format_query") is False
+
+    def test_blocked_when_results_focused(self):
+        sm = UIStateMachine()
+        ctx = make_context(
+            focus="results",
+            leader_pending=True,
+            leader_menu="leader",
+        )
+        assert sm.check_action(ctx, "leader_format_query") is False
+
+    def test_help_lists_space_p_format_binding(self):
+        sm = UIStateMachine()
+        query_help = next(
+            section for section in sm.generate_help_sections() if section.id == "query_normal"
+        )
+
+        assert any(
+            item.kind == "binding"
+            and item.key == "<space>p"
+            and item.description == "Format query"
+            for item in query_help.items
+        )

tests/ui/test_query_formatting.py

@@ -0,0 +1,84 @@
+"""Pilot-driven query formatting tests."""
+
+from __future__ import annotations
+
+import pytest
+
+from sqlit.domains.shell.app.main import SSMSTUI
+
+from .mocks import MockConnectionStore, MockSettingsStore, build_test_services
+
+
+def _make_app() -> SSMSTUI:
+    return SSMSTUI(
+        services=build_test_services(
+            connection_store=MockConnectionStore(),
+            settings_store=MockSettingsStore({"theme": "tokyo-night"}),
+        )
+    )
+
+
+@pytest.mark.asyncio
+async def test_space_p_formats_query_through_real_key_dispatch() -> None:
+    app = _make_app()
+
+    async with app.run_test(size=(100, 35)) as pilot:
+        await pilot.pause()
+        app.action_focus_query()
+        app.query_input.text = "select id,name from users where active=true"
+        await pilot.press("space", "p")
+        await pilot.pause()
+
+        assert app.query_input.text == """SELECT id,
+       name
+FROM users
+WHERE active = TRUE"""
+
+
+@pytest.mark.asyncio
+async def test_format_action_is_undoable() -> None:
+    app = _make_app()
+    original = "select id,name from users"
+
+    async with app.run_test(size=(100, 35)) as pilot:
+        await pilot.pause()
+        app.action_focus_query()
+        app.query_input.text = original
+        await pilot.press("space", "p")
+        await pilot.pause()
+        await pilot.press("u")
+        await pilot.pause()
+
+        assert app.query_input.text == original
+
+
+@pytest.mark.asyncio
+async def test_format_action_preserves_cursor_token() -> None:
+    app = _make_app()
+    query = "select id,name from users where active=true"
+
+    async with app.run_test(size=(100, 35)) as pilot:
+        await pilot.pause()
+        app.action_focus_query()
+        app.query_input.text = query
+        app.query_input.cursor_location = (0, query.index("users") + len("users"))
+        await pilot.press("space", "p")
+        await pilot.pause()
+
+        row, column = app.query_input.cursor_location
+        assert app.query_input.text.splitlines()[row][:column].endswith("users")
+
+
+@pytest.mark.asyncio
+async def test_empty_query_is_not_changed() -> None:
+    app = _make_app()
+
+    async with app.run_test(size=(100, 35)) as pilot:
+        await pilot.pause()
+        app.action_focus_query()
+        app.query_input.text = ""
+        await pilot.press("space", "p")
+        await pilot.pause()
+
+        assert app.query_input.text == ""
+        assert not app._get_undo_history().can_undo()