Refactoring - query optimizations for large results - UX feedback improvements

Author: Maxteabag

README.md

@@ -16,18 +16,16 @@ A lightweight TUI for people who just want to run some queries fast.
 
 ## Features
 
-- **Multi-database support**: SQL Server, PostgreSQL, MySQL, SQLite, MariaDB, Oracle, DuckDB, and CockroachDB
-- **SSH tunnel support**: Connect to remote databases through SSH with password or key-based authentication
-- **Query history**: Automatically saves queries per connection, searchable and sortable
-- Fast and intuitive keyboard only control
-- Context based help (no need to memorize tons of hot-keys)
+- **Connection manager UI** - Save connections, switch between databases without CLI args
+- **Just run `sqlit`** - No CLI config needed, pick a connection and go
+- **Multi-database out of the box** - SQL Server, PostgreSQL, MySQL, SQLite, MariaDB, Oracle, DuckDB, CockroachDB - no adapters to install
+- **SSH tunnels built-in** - Connect to remote databases securely with password or key auth
+- **Vim-style editing** - Modal editing for terminal purists
+- **Query history** - Automatically saves queries per connection, searchable and sortable
+- Context-aware help (no need to memorize keybindings)
 - Browse databases, tables, views, and stored procedures
-- Execute SQL queries with syntax highlighting
-- Vim-style query editing
 - SQL autocomplete for tables, columns, and procedures
-- Multiple authentication methods (Windows, SQL Server, Entra ID)
-- Save and manage connections
-- Responsive terminal UI
+- Multiple auth methods (Windows, SQL Server, Entra ID)
 - CLI mode for scripting and AI agents
 - Themes (Tokyo Night, Nord, and more)
 - Auto-detects and installs ODBC drivers (SQL Server)

demo-connect.gif

@@ -199,6 +199,8 @@ class SSMSTUI(
         Binding("Y", "copy_row", "Copy row", show=False),
         Binding("a", "copy_results", "Copy results", show=False),
         Binding("ctrl+c", "cancel_operation", "Cancel", show=False),
+        Binding("N", "show_notifications", "Notifications", show=False),
+        Binding("d", "dismiss_notification", "Dismiss", show=False),
     ]
 
     def __init__(self):
@@ -210,6 +212,8 @@ def __init__(self):
         self.current_ssh_tunnel: Any | None = None
         self.vim_mode: VimMode = VimMode.NORMAL
         self._expanded_paths: set[str] = set()
+        self._loading_nodes: set[str] = set()
+        self._session: Any | None = None
         self._schema_cache: dict = {
             "tables": [],
             "views": [],
@@ -226,9 +230,14 @@ def __init__(self):
         self._last_result_row_count: int = 0
         self._internal_clipboard: str = ""
         self._fullscreen_mode: str = "none"
-        self._connection_health: dict[str, bool] = {}
+        self._last_notification: str = ""
+        self._last_notification_severity: str = "information"
+        self._last_notification_time: str = ""
+        self._notification_timer = None
+        self._notification_history: list = []
         self._query_worker = None
         self._query_executing: bool = False
+        self._cancellable_query: Any | None = None
         self._spinner_index: int = 0
         self._spinner_timer = None
         # Schema indexing state
@@ -237,6 +246,7 @@ def __init__(self):
         self._schema_spinner_index: int = 0
         self._schema_spinner_timer = None
         self._table_metadata: dict = {}
+        self._columns_loading: set[str] = set()
 
     def push_screen(self, screen, callback=None, wait_for_dismiss: bool = False):
         """Override push_screen to hide footer when showing modal dialogs."""

demo-history.gif

@@ -104,6 +104,13 @@ def main() -> int:
         choices=["table", "csv", "json"],
         help="Output format (default: table)",
     )
+    query_parser.add_argument(
+        "--limit",
+        "-l",
+        type=int,
+        default=1000,
+        help="Maximum rows to fetch (default: 1000, use 0 for unlimited)",
+    )
 
     args = parser.parse_args()
 

sqlit/app.py

@@ -5,8 +5,8 @@
 import csv
 import json
 import sys
+from typing import TYPE_CHECKING, Callable
 
-from .adapters import create_ssh_tunnel, get_adapter
 from .config import (
     AUTH_TYPE_LABELS,
     AuthType,
@@ -16,6 +16,10 @@
     load_connections,
     save_connections,
 )
+from .services import ConnectionSession, QueryResult, QueryService
+
+if TYPE_CHECKING:
+    from .services import HistoryStoreProtocol
 
 
 def cmd_connection_list(args) -> int:
@@ -219,8 +223,98 @@ def cmd_connection_delete(args) -> int:
     return 0
 
 
-def cmd_query(args) -> int:
-    """Execute a SQL query against a connection."""
+def _stream_csv_output(cursor, columns: list[str]) -> int:
+    """Stream CSV output from cursor using fetchmany."""
+    writer = csv.writer(sys.stdout)
+    writer.writerow(columns)
+    row_count = 0
+    batch_size = 1000
+    while True:
+        rows = cursor.fetchmany(batch_size)
+        if not rows:
+            break
+        for row in rows:
+            writer.writerow(str(val) if val is not None else "" for val in row)
+            row_count += 1
+    return row_count
+
+
+def _stream_json_output(cursor, columns: list[str]) -> int:
+    """Stream JSON output from cursor using fetchmany (JSON array format)."""
+    print("[")
+    first = True
+    row_count = 0
+    batch_size = 1000
+    while True:
+        rows = cursor.fetchmany(batch_size)
+        if not rows:
+            break
+        for row in rows:
+            if not first:
+                print(",")
+            first = False
+            obj = dict(zip(columns, [val if val is not None else None for val in row]))
+            print(json.dumps(obj, default=str), end="")
+            row_count += 1
+    print("\n]")
+    return row_count
+
+
+def _output_table(columns: list[str], rows: list[tuple], truncated: bool) -> None:
+    """Output query results in table format with optimized width calculation."""
+    MAX_COL_WIDTH = 50  # Cap column width to avoid excessive line length
+
+    # Calculate column widths (only scan first 100 rows for performance)
+    col_widths = [min(len(col), MAX_COL_WIDTH) for col in columns]
+    for row in rows[:100]:
+        for i, val in enumerate(row):
+            val_str = str(val) if val is not None else "NULL"
+            col_widths[i] = min(MAX_COL_WIDTH, max(col_widths[i], len(val_str)))
+
+    # Print header
+    header_parts = []
+    for i, col in enumerate(columns):
+        col_display = col[:col_widths[i]] if len(col) > col_widths[i] else col
+        header_parts.append(col_display.ljust(col_widths[i]))
+    header = " | ".join(header_parts)
+    print(header)
+    print("-" * len(header))
+
+    # Print rows
+    for row in rows:
+        row_parts = []
+        for i, val in enumerate(row):
+            val_str = str(val) if val is not None else "NULL"
+            if len(val_str) > col_widths[i]:
+                val_str = val_str[: col_widths[i] - 2] + ".."
+            row_parts.append(val_str.ljust(col_widths[i]))
+        print(" | ".join(row_parts))
+
+    # Print count with truncation notice
+    if truncated:
+        print(f"\n({len(rows)} rows shown, results truncated)")
+    else:
+        print(f"\n({len(rows)} row(s) returned)")
+
+
+def cmd_query(
+    args,
+    *,
+    session_factory: Callable[[ConnectionConfig], ConnectionSession] | None = None,
+    query_service: QueryService | None = None,
+) -> int:
+    """Execute a SQL query against a connection.
+
+    Args:
+        args: Parsed command-line arguments.
+        session_factory: Optional factory for creating ConnectionSession.
+            Defaults to ConnectionSession.create. Useful for testing.
+        query_service: Optional QueryService instance.
+            Defaults to a new QueryService(). Useful for testing.
+
+    Returns:
+        Exit code (0 for success, 1 for error).
+    """
     connections = load_connections()
 
     config = None
@@ -253,86 +347,84 @@ def cmd_query(args) -> int:
         print("Error: Either --query or --file must be provided.")
         return 1
 
-    tunnel = None
+    # Determine row limit (0 means unlimited)
+    max_rows = args.limit if args.limit > 0 else None
+
+    # Use injected or default factories
+    create_session = session_factory or ConnectionSession.create
+    service = query_service or QueryService()
+
     try:
-        from dataclasses import replace
-
-        # Create SSH tunnel if enabled
-        tunnel, host, port = create_ssh_tunnel(config)
-        if tunnel:
-            connect_config = replace(config, server=host, port=str(port))
-        else:
-            connect_config = config
-
-        adapter = get_adapter(config.db_type)
-        db_conn = adapter.connect(connect_config)
-
-        # Detect query type to avoid executing non-SELECT statements twice
-        query_type = query.strip().upper().split()[0] if query.strip() else ""
-        is_select_query = query_type in ("SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA")
-
-        if is_select_query:
-            columns, rows, _truncated = adapter.execute_query(db_conn, query)
-        else:
-            columns, rows = [], []
-
-        if columns:
-            if args.format == "csv":
-                # Use proper CSV writer for correct quoting/escaping
-                writer = csv.writer(sys.stdout)
-                writer.writerow(columns)
-                for row in rows:
-                    writer.writerow(str(val) if val is not None else "" for val in row)
-            elif args.format == "json":
-                result = []
-                for row in rows:
-                    result.append(
-                        dict(
-                            zip(
-                                columns,
-                                [val if val is not None else None for val in row],
-                            )
-                        )
-                    )
-                print(json.dumps(result, indent=2, default=str))
+        # Use ConnectionSession for automatic resource cleanup
+        with create_session(config) as session:
+            # For unlimited streaming output (CSV/JSON only), use direct cursor access
+            from .services.query import is_select_query
+
+            if max_rows is None and args.format in ("csv", "json") and is_select_query(query):
+                # Stream directly from cursor for unlimited CSV/JSON
+                cursor = session.connection.cursor()
+                cursor.execute(query)
+
+                if not cursor.description:
+                    print("Query executed successfully (no results)")
+                    return 0
+
+                columns = [col[0] for col in cursor.description]
+
+                if args.format == "csv":
+                    row_count = _stream_csv_output(cursor, columns)
+                else:
+                    row_count = _stream_json_output(cursor, columns)
+
+                # Save to history
+                service._save_to_history(config.name, query)
+                print(f"\n({row_count} row(s) returned)", file=sys.stderr)
+                return 0
+
+            # Standard execution with QueryService (with row limit)
+            result = service.execute(
+                connection=session.connection,
+                adapter=session.adapter,
+                query=query,
+                config=config,
+                max_rows=max_rows,
+                save_to_history=True,
+            )
+
+            if isinstance(result, QueryResult):
+                columns = result.columns
+                rows = result.rows
+
+                if args.format == "csv":
+                    writer = csv.writer(sys.stdout)
+                    writer.writerow(columns)
+                    for row in rows:
+                        writer.writerow(str(val) if val is not None else "" for val in row)
+                    if result.truncated:
+                        print(f"\n({len(rows)} rows shown, results truncated)", file=sys.stderr)
+                    else:
+                        print(f"\n({len(rows)} row(s) returned)", file=sys.stderr)
+                elif args.format == "json":
+                    json_result = [
+                        dict(zip(columns, [val if val is not None else None for val in row]))
+                        for row in rows
+                    ]
+                    print(json.dumps(json_result, indent=2, default=str))
+                    if result.truncated:
+                        print(f"\n({len(rows)} rows shown, results truncated)", file=sys.stderr)
+                    else:
+                        print(f"\n({len(rows)} row(s) returned)", file=sys.stderr)
+                else:
+                    _output_table(columns, rows, result.truncated)
             else:
-                col_widths = [len(col) for col in columns]
-                for row in rows:
-                    for i, val in enumerate(row):
-                        col_widths[i] = max(
-                            col_widths[i], len(str(val) if val is not None else "NULL")
-                        )
-
-                header = " | ".join(
-                    col.ljust(col_widths[i]) for i, col in enumerate(columns)
-                )
-                print(header)
-                print("-" * len(header))
-
-                for row in rows:
-                    row_str = " | ".join(
-                        (str(val) if val is not None else "NULL").ljust(col_widths[i])
-                        for i, val in enumerate(row)
-                    )
-                    print(row_str)
-
-            print(f"\n({len(rows)} row(s) returned)")
-        else:
-            affected = adapter.execute_non_query(db_conn, query)
-            print(f"Query executed successfully. Rows affected: {affected}")
-
-        db_conn.close()
-        if tunnel:
-            tunnel.stop()
-        return 0
+                # NonQueryResult
+                print(f"Query executed successfully. Rows affected: {result.rows_affected}")
+
+            return 0
 
     except ImportError as e:
         print(f"Error: Required module not installed: {e}")
-        if tunnel:
-            tunnel.stop()
         return 1
     except Exception as e:
         print(f"Error: {e}")
-        if tunnel:
-            tunnel.stop()
         return 1

sqlit/cli.py

@@ -124,7 +124,22 @@ def get_auth_type(self) -> AuthType:
             return AuthType.SQL_SERVER
 
     def get_connection_string(self) -> str:
-        """Build the connection string for SQL Server."""
+        """Build the connection string for SQL Server.
+
+        .. deprecated::
+            This method is deprecated. Connection string building is now
+            handled internally by SQLServerAdapter._build_connection_string().
+            Use SQLServerAdapter.connect() directly instead.
+        """
+        import warnings
+
+        warnings.warn(
+            "ConnectionConfig.get_connection_string() is deprecated. "
+            "Connection string building is now handled internally by SQLServerAdapter.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
         if self.db_type != "mssql":
             raise ValueError("get_connection_string() is only for SQL Server connections")