Coverage for src/ui/tui.py: 58%

893 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 23:16 -0700

1""" 

2Main TUI interface for CHUCK AI. 

3""" 

4 

5import os 

6import shlex 

7from typing import List, Dict, Any 

8import logging 

9 

10from rich.console import Console 

11import traceback 

12 

13# Rich imports for TUI rendering 

14from rich.panel import Panel 

15 

16# Prompt Toolkit imports for enhanced CLI experience 

17from prompt_toolkit import PromptSession 

18from prompt_toolkit.history import FileHistory 

19from prompt_toolkit.auto_suggest import AutoSuggestFromHistory 

20from prompt_toolkit.completion import WordCompleter 

21from prompt_toolkit.formatted_text import HTML 

22from prompt_toolkit.styles import Style 

23from prompt_toolkit.key_binding import KeyBindings 

24 

25from src.ui.ascii_art import display_welcome_screen 

26 

27from src.ui.theme import ( 

28 WARNING, 

29 INFO, 

30 INFO_STYLE, 

31 SUCCESS_STYLE, 

32 ERROR_STYLE, 

33 WARNING_STYLE, 

34 DIALOG_BORDER, 

35 TABLE_TITLE_STYLE, 

36) 

37 

38from src.service import ChuckService 

39from src.commands.base import CommandResult 

40from src.command_registry import get_command 

41from src.config import get_active_model 

42 

43# Import the interactive context manager 

44from src.interactive_context import InteractiveContext 

45 

46# Global reference to TUI instance for service access 

47_tui_instance = None 

48 

49 

50def get_chuck_service(): 

51 """Get the global ChuckService instance.""" 

52 if _tui_instance is None: 

53 return None 

54 return _tui_instance.get_service() 

55 

56 

57def set_chuck_service(service): 

58 """Set a new global ChuckService instance.""" 

59 if _tui_instance is None: 

60 return False 

61 return _tui_instance.set_service(service) 

62 

63 

64def get_console(): 

65 """Get the global TUI console instance.""" 

66 if _tui_instance is None: 

67 # Fallback to a default console if TUI not available 

68 from rich.console import Console 

69 

70 return Console() 

71 return _tui_instance.console 

72 

73 

74class ChuckTUI: 

75 """ 

76 Main TUI interface for CHUCK AI. 

77 Handles user interaction, execution via ChuckService, and result display. 

78 """ 

79 

80 def __init__(self, no_color=False): 

81 """Initialize the CHUCK AI TUI.""" 

82 self.console = Console(force_terminal=not no_color, no_color=no_color) 

83 self.service = ChuckService() 

84 self.running = True 

85 self.debug = False # Debug state 

86 self.no_color = no_color 

87 

88 # Register this instance as the global TUI instance 

89 # This allows other modules to access the service instance 

90 global _tui_instance 

91 _tui_instance = self 

92 

93 def get_service(self): 

94 """Get the current ChuckService instance.""" 

95 return self.service 

96 

97 def set_service(self, service): 

98 """Set a new ChuckService instance.""" 

99 self.service = service 

100 return True 

101 

102 def _get_available_commands(self) -> List[str]: 

103 """ 

104 Get a list of available commands for autocompletion. 

105 """ 

106 # Built-in commands always available 

107 builtin_commands = ["/exit", "/quit", "/help", "/debug"] 

108 

109 # Add service commands from the command registry 

110 service_commands = [] 

111 try: 

112 # Use the command registry instead of hardcoding 

113 from src.command_registry import TUI_COMMAND_MAP 

114 

115 service_commands = list(TUI_COMMAND_MAP.keys()) 

116 except Exception as e: 

117 if self.debug: 

118 self.console.print(f"[dim]Error getting commands: {str(e)}[/dim]") 

119 

120 # Combine all commands and remove duplicates 

121 all_commands = builtin_commands + service_commands 

122 return sorted(list(set(all_commands))) 

123 

124 def _check_first_run(self) -> bool: 

125 """ 

126 Check if this is the first run of the application and configuration is needed. 

127 

128 Returns: 

129 True if setup wizard should be launched, False otherwise 

130 """ 

131 from src.config import get_config_manager 

132 from pathlib import Path 

133 

134 config_manager = get_config_manager() 

135 

136 if config_manager.needs_setup(): 

137 if not Path(config_manager.config_path).exists(): 

138 self.console.print( 

139 f"[{WARNING_STYLE}]First time running Chuck! Starting setup wizard...[/{WARNING_STYLE}]" 

140 ) 

141 return True 

142 else: 

143 self.console.print( 

144 f"[{WARNING_STYLE}]Some configuration settings are missing. Starting setup wizard...[/{WARNING_STYLE}]" 

145 ) 

146 return True 

147 

148 return False 

149 

150 def run(self) -> None: 

151 """ 

152 Run the CHUCK AI TUI interface with enhanced prompt-toolkit interface. 

153 """ 

154 # Clear the screen (skip in no-color mode for better integration test compatibility) 

155 if not self.no_color: 

156 os.system("cls" if os.name == "nt" else "clear") 

157 

158 # Display welcome screen 

159 display_welcome_screen(self.console) 

160 

161 # Initialize the interactive context 

162 interactive_context = InteractiveContext() 

163 history_file = os.path.expanduser("~/.chuck_history") 

164 

165 # Get available commands from service for autocomplete 

166 commands = self._get_available_commands() 

167 

168 # Set up command completer with slash prefix 

169 command_completer = WordCompleter( 

170 commands, ignore_case=True, match_middle=True, sentence=True 

171 ) 

172 

173 # Custom prompt styles to match our Rich formatting 

174 # Use colors that are consistent with our theme, but respect no_color setting 

175 # Note: prompt-toolkit requires specific ANSI color names 

176 if self.no_color: 

177 style = Style.from_dict( 

178 { 

179 "prompt": "", # No styling in no-color mode 

180 "interactive-prompt": "", # No styling in no-color mode 

181 } 

182 ) 

183 else: 

184 style = Style.from_dict( 

185 { 

186 "prompt": "ansicyan bold", # Matches TABLE_TITLE_STYLE 

187 "interactive-prompt": "ansiblue bold", # Matches DIALOG_BORDER 

188 } 

189 ) 

190 

191 # Key bindings to accept history suggestions with Tab 

192 bindings = KeyBindings() 

193 

194 @bindings.add("tab") 

195 def _(event): 

196 buff = event.current_buffer 

197 if buff.suggestion: 

198 buff.insert_text(buff.suggestion.text) 

199 else: 

200 if buff.complete_state: 

201 buff.complete_next() 

202 else: 

203 buff.start_completion(select_first=True) 

204 

205 # Create the prompt session with history and completion 

206 # Disable syntax highlighting in no-color mode 

207 session = PromptSession( 

208 history=FileHistory(history_file), 

209 auto_suggest=AutoSuggestFromHistory(), 

210 completer=command_completer, 

211 complete_while_typing=True, 

212 style=style, 

213 lexer=None, 

214 include_default_pygments_style=False, 

215 key_bindings=bindings, 

216 ) 

217 

218 # Check if this is the first run and we need to set up 

219 if self._check_first_run(): 

220 # Start the setup wizard in interactive mode 

221 self.console.print("[bold]Starting Chuck setup wizard...[/bold]") 

222 # Just call the setup command - it will handle its own context 

223 result = self.service.execute_command("/setup") 

224 

225 # Main TUI application loop 

226 while self.running: 

227 try: 

228 # Check if we're in interactive mode 

229 if interactive_context.is_in_interactive_mode(): 

230 current_cmd = interactive_context.current_command 

231 

232 # Use prompt toolkit with interactive styling 

233 prompt_message = HTML( 

234 "<interactive-prompt>chuck (interactive) ></interactive-prompt> " 

235 ) 

236 

237 # Determine if we should hide input (e.g., when entering tokens) 

238 hide_input = False 

239 if current_cmd == "/setup": 

240 ctx = interactive_context.get_context_data("/setup") 

241 # We're on the token input step of the setup wizard 

242 step = ctx.get("current_step") 

243 # Only hide input if we're specifically on the token input step AND have a workspace URL 

244 if step == "token_input" and ctx.get("workspace_url"): 

245 hide_input = True 

246 

247 user_input = session.prompt( 

248 prompt_message, 

249 is_password=hide_input, 

250 enable_history_search=not hide_input, 

251 ).strip() 

252 

253 # Allow escaping from interactive mode 

254 if user_input.lower() in ["cancel", "exit", "quit"]: 

255 interactive_context.clear_active_context(current_cmd) 

256 self.console.print( 

257 f"[{WARNING_STYLE}]Exited interactive mode[/{WARNING_STYLE}]" 

258 ) 

259 continue 

260 

261 # Process the interactive input 

262 result = self.service.execute_command( 

263 current_cmd, interactive_input=user_input 

264 ) 

265 

266 # Process the result even in interactive mode 

267 if result: 

268 # Check for special exit_interactive message 

269 if ( 

270 not result.success 

271 and result.message 

272 and result.message.startswith("exit_interactive:") 

273 ): 

274 # Extract the actual message after the prefix 

275 actual_message = ( 

276 result.message.split(":", 1)[1] 

277 if ":" in result.message 

278 else result.message 

279 ) 

280 interactive_context.clear_active_context(current_cmd) 

281 self.console.print( 

282 f"[{INFO_STYLE}]{actual_message}[/{INFO_STYLE}]" 

283 ) 

284 continue 

285 

286 # Check if the command rendered content and we should not show a message 

287 if ( 

288 result.success 

289 and isinstance(result.data, dict) 

290 and result.data.get("rendered_next_step") 

291 ): 

292 # Content was rendered, just continue to next prompt without additional output 

293 continue 

294 

295 # If the command is no longer in interactive mode, it means it's complete 

296 if ( 

297 not interactive_context.is_in_interactive_mode() 

298 or interactive_context.current_command != current_cmd 

299 ): 

300 # Command completed its interactive flow, show final result 

301 if result.success and result.message: 

302 self.console.print( 

303 f"[bold green]{result.message}[/bold green]" 

304 ) 

305 elif not result.success: 

306 self._display_error(result) 

307 else: 

308 # Regular command mode - use prompt toolkit 

309 prompt_message = HTML("<prompt>chuck ></prompt> ") 

310 command = session.prompt(prompt_message).strip() 

311 

312 # Skip empty commands 

313 if not command: 

314 continue 

315 

316 # Process user command 

317 self._process_command(command) 

318 

319 except KeyboardInterrupt: 

320 # Handle Ctrl+C 

321 self.console.print( 

322 f"\n[{WARNING_STYLE}]Interrupted by user. Type 'exit' to quit.[/{WARNING_STYLE}]" 

323 ) 

324 except Exception as e: 

325 # Import at the top of the method to avoid scoping issues 

326 from src.exceptions import PaginationCancelled 

327 

328 if isinstance(e, PaginationCancelled): 

329 # Handle pagination cancellation silently - just return to prompt 

330 pass 

331 else: 

332 raise # Re-raise other exceptions 

333 except EOFError: 

334 # Handle Ctrl+D 

335 self.console.print( 

336 f"\n[{WARNING_STYLE}]Thank you for using chuck![/{WARNING_STYLE}]" 

337 ) 

338 break 

339 except Exception as e: 

340 # Handle other exceptions 

341 self.console.print( 

342 f"[{ERROR_STYLE}]Unexpected Error: {str(e)}[/{ERROR_STYLE}]" 

343 ) 

344 # Print stack trace in debug mode 

345 if self.debug: 

346 self.console.print("[dim]" + traceback.format_exc() + "[/dim]") 

347 

348 def _needs_shlex_parsing(self, command: str) -> bool: 

349 """Determine if command needs shlex parsing (has quotes or flags).""" 

350 # Never use shlex for agent commands - they need simple splitting 

351 # to preserve natural language as-is 

352 if command.startswith("/agent "): 

353 return False 

354 

355 # Use shlex parsing for commands that clearly need it: 

356 # 1. Flag-style arguments (--flag) 

357 # 2. Intentional quoted strings (balanced quotes) 

358 

359 # Always use shlex for flag-style arguments 

360 if "--" in command: 

361 return True 

362 

363 # For quoted strings, only use shlex if quotes appear balanced 

364 if '"' in command: 

365 double_quotes = command.count('"') 

366 if double_quotes >= 2 and double_quotes % 2 == 0: 

367 return True 

368 

369 # For single quotes, be more careful - avoid contractions like "let's" 

370 if "'" in command: 

371 single_quotes = command.count("'") 

372 # Only use shlex if we have multiple balanced single quotes 

373 # and they don't appear to be contractions 

374 if single_quotes >= 2 and single_quotes % 2 == 0: 

375 # Additional check: make sure quotes aren't part of contractions 

376 # This is a simple heuristic - look for patterns like "'s " or "'t " 

377 if not ( 

378 "'s " in command 

379 or "'t " in command 

380 or "'ll " in command 

381 or "'re " in command 

382 or "'ve " in command 

383 or "'d " in command 

384 ): 

385 return True 

386 

387 return False 

388 

389 def _process_command(self, command): 

390 """Process a user command in the TUI interface.""" 

391 interactive_context = InteractiveContext() 

392 

393 # Handle built-in TUI commands 

394 if command.lower() in ["/exit", "/quit", "exit", "quit"]: 

395 self.running = False 

396 self.console.print( 

397 f"[{WARNING_STYLE}]Exiting Chuck AI...[/{WARNING_STYLE}]" 

398 ) 

399 return 

400 

401 # Process command with agent if no slash prefix 

402 if not command.startswith("/"): 

403 self.console.print("[teal]Thinking...[/teal]") 

404 command = f"/agent {command}" # Default to agent if no slash 

405 elif command.startswith("/ask "): 

406 # Show thinking message for /ask commands too 

407 self.console.print("[teal]Thinking...[/teal]") 

408 

409 # Split command into parts for service layer 

410 # Use shlex for commands with quotes or flags, simple split for natural language 

411 if self._needs_shlex_parsing(command): 

412 try: 

413 parts = shlex.split(command) 

414 except ValueError as e: 

415 self.console.print(f"[red]Error parsing command: {e}[/red]") 

416 return 

417 else: 

418 parts = command.split() 

419 

420 cmd = parts[0].lower() 

421 args = parts[1:] 

422 

423 # Process special commands 

424 if cmd == "/debug": 

425 self._handle_debug(args) 

426 return 

427 

428 # Execute command via service layer 

429 # Pass display callback for agent commands to show tool outputs immediately 

430 if cmd in ["/agent", "/ask"]: 

431 result = self.service.execute_command( 

432 cmd, *args, tool_output_callback=self.display_tool_output 

433 ) 

434 else: 

435 result = self.service.execute_command(cmd, *args) 

436 

437 if not result: 

438 return 

439 

440 # Skip result processing if entering interactive mode 

441 # Ensure valid commands initiated outside of agent tools are not skipped 

442 if interactive_context.is_in_interactive_mode() and cmd not in [ 

443 "/agent", 

444 "/ask", 

445 ]: 

446 return 

447 

448 # Process command result 

449 self._process_command_result(cmd, result) 

450 

451 def _process_command_result(self, cmd, result): 

452 """Process a command result and display it appropriately in the TUI.""" 

453 if result.success: 

454 # Display success message if available 

455 if result.message: 

456 # Check for specific data types to format messages differently 

457 if ( 

458 cmd in ["/agent", "/ask"] 

459 and isinstance(result.data, dict) 

460 and "response" in result.data 

461 ): 

462 # Agent response is handled below, just print success message if any 

463 if ( 

464 result.message != result.data["response"] 

465 ): # Avoid duplicate printing 

466 self.console.print(f"[bold green]{result.message}[/bold green]") 

467 else: 

468 self.console.print(f"[bold green]{result.message}[/bold green]") 

469 

470 # Skip if no data to display 

471 if not result.data: 

472 return 

473 

474 # Specialized display for different commands 

475 if cmd in ["/catalogs", "/search_catalogs", "/list-catalogs"]: 

476 self._display_catalogs(result.data) 

477 elif cmd in ["/schemas", "/search_schemas", "/list-schemas"]: 

478 self._display_schemas(result.data) 

479 elif cmd in ["/tables", "/search_tables", "/list-tables"]: 

480 self._display_tables(result.data) 

481 elif cmd in ["/catalog", "/catalog-details"]: 

482 self._display_catalog_details(result.data) 

483 elif cmd in ["/schema", "/schema-details"]: 

484 self._display_schema_details(result.data) 

485 elif cmd == "/models": 

486 self._display_models(result.data) 

487 elif cmd == "/list-models": 

488 self._display_detailed_models(result.data) 

489 elif cmd in ["/warehouses", "/list-warehouses"]: 

490 self._display_warehouses(result.data) 

491 elif cmd in ["/volumes", "/list-volumes"]: 

492 self._display_volumes(result.data) 

493 elif cmd in ["/table", "/show_table", "/table-details"]: 

494 self._display_table_details(result.data) 

495 elif cmd in ["/run-sql", "/sql"]: 

496 self._display_sql_results(result.data) 

497 elif cmd == "/scan-pii": 

498 self._display_pii_scan_results(result.data) 

499 elif cmd == "/status": 

500 self._display_status(result.data) 

501 elif cmd == "/auth" and "permissions" in result.data: 

502 self._display_permissions(result.data["permissions"]) 

503 elif cmd == "/usage": 

504 # For the usage command, we just display the message 

505 if result.message: 

506 self._display_usage(result.message) 

507 elif ( 

508 cmd == "/help" 

509 and isinstance(result.data, dict) 

510 and "help_text" in result.data 

511 ): 

512 self.console.print( 

513 Panel( 

514 result.data["help_text"], 

515 title="CHUCK AI Help", 

516 border_style="cyan", 

517 ) 

518 ) 

519 elif ( 

520 cmd in ["/agent", "/ask"] 

521 and isinstance(result.data, dict) 

522 and "response" in result.data 

523 and result.data[ 

524 "response" 

525 ].strip() # Only show if response is not empty 

526 ): 

527 # Agent response - print directly or format nicely 

528 self.console.print( 

529 Panel( 

530 result.data["response"], 

531 title="Agent Response", 

532 border_style=DIALOG_BORDER, 

533 ) 

534 ) 

535 else: 

536 # Display error 

537 self._display_error(result) 

538 

539 def _handle_debug(self, args: List[str]) -> None: 

540 """Toggle debug mode.""" 

541 if not args: 

542 self.debug = not self.debug 

543 elif args[0].lower() in ("on", "true", "1", "yes"): 

544 self.debug = True 

545 elif args[0].lower() in ("off", "false", "0", "no"): 

546 self.debug = False 

547 else: 

548 self.console.print(f"[{ERROR_STYLE}]Invalid debug option.[/{ERROR_STYLE}]") 

549 self.console.print("Usage: /debug [on|off]") 

550 return 

551 

552 status = "ON" if self.debug else "OFF" 

553 self.console.print( 

554 f"[{SUCCESS_STYLE}]Debug mode is now {status}[/{SUCCESS_STYLE}]" 

555 ) 

556 

557 def display_tool_output(self, tool_name: str, tool_result: Dict[str, Any]) -> None: 

558 """Display tool output immediately during agent execution.""" 

559 try: 

560 # Get command definition to check display type 

561 

562 command_def = get_command(tool_name) 

563 

564 # Use the command's agent_display setting, defaulting to "condensed" 

565 display_type = "condensed" 

566 if command_def: 

567 display_type = getattr(command_def, "agent_display", "condensed") 

568 

569 # Route based on display type 

570 if display_type == "condensed": 

571 self._display_condensed_tool_output(tool_name, tool_result) 

572 else: 

573 # Full display - use existing detailed display methods 

574 self._display_full_tool_output(tool_name, tool_result) 

575 

576 except Exception as e: 

577 # Handle pagination cancellation specially - let it bubble up 

578 from src.exceptions import PaginationCancelled 

579 

580 if isinstance(e, PaginationCancelled): 

581 raise # Re-raise to bubble up to main TUI loop 

582 

583 # Don't let other display errors break agent execution 

584 logging.warning(f"Failed to display tool output for {tool_name}: {e}") 

585 # Show a simple notification that output was attempted 

586 self.console.print(f"[dim][Tool: {tool_name} executed][/dim]") 

587 

588 def _display_full_tool_output( 

589 self, tool_name: str, tool_result: Dict[str, Any] 

590 ) -> None: 

591 """Display full detailed tool output (existing behavior).""" 

592 # Map tool names to their display methods 

593 # This reuses existing display logic to maintain consistency 

594 if tool_name in ["list-catalogs", "list_catalogs", "catalogs"]: 

595 self._display_catalogs(tool_result) 

596 elif tool_name in ["list-schemas", "list_schemas", "schemas"]: 

597 self._display_schemas(tool_result) 

598 elif tool_name in ["list-tables", "list_tables", "tables"]: 

599 self._display_tables(tool_result) 

600 elif tool_name in ["get_catalog_details", "catalog"]: 

601 self._display_catalog_details(tool_result) 

602 elif tool_name in ["get_schema_details", "schema"]: 

603 self._display_schema_details(tool_result) 

604 elif tool_name in ["detailed-models", "list-models", "list_models", "models"]: 

605 if "models" in tool_result: 

606 self._display_detailed_models(tool_result) 

607 else: 

608 self._display_models(tool_result) 

609 elif tool_name in ["list-warehouses", "list_warehouses", "warehouses"]: 

610 self._display_warehouses(tool_result) 

611 elif tool_name in ["list-volumes", "list_volumes", "volumes"]: 

612 self._display_volumes(tool_result) 

613 elif tool_name in ["get_table_info", "table", "show_table"]: 

614 self._display_table_details(tool_result) 

615 elif tool_name in ["scan_schema_for_pii", "scan_pii"]: 

616 self._display_pii_scan_results(tool_result) 

617 elif tool_name == "get_status": 

618 self._display_status(tool_result) 

619 elif tool_name == "run-sql": 

620 self._display_sql_results_formatted(tool_result) 

621 else: 

622 # For unknown tools, display a generic panel with the data 

623 from rich.panel import Panel 

624 import json 

625 

626 # Try to format as JSON for readability 

627 try: 

628 formatted_data = json.dumps(tool_result, indent=2) 

629 self.console.print( 

630 Panel( 

631 formatted_data, 

632 title=f"Tool Output: {tool_name}", 

633 border_style=DIALOG_BORDER, 

634 ) 

635 ) 

636 except (TypeError, ValueError): 

637 # If JSON serialization fails, display as string 

638 self.console.print( 

639 Panel( 

640 str(tool_result), 

641 title=f"Tool Output: {tool_name}", 

642 border_style=DIALOG_BORDER, 

643 ) 

644 ) 

645 

646 def _display_condensed_tool_output( 

647 self, tool_name: str, tool_result: Dict[str, Any] 

648 ) -> None: 

649 """Display condensed tool output with status and key metrics.""" 

650 # Get friendly action name if available 

651 

652 command_def = get_command(tool_name) 

653 friendly_name = tool_name 

654 if command_def and getattr(command_def, "condensed_action", None): 

655 friendly_name = command_def.condensed_action 

656 

657 # Extract key information based on tool type 

658 status_line = f"[dim cyan]→[/dim cyan] {friendly_name}" 

659 metrics = [] 

660 

661 # Extract meaningful metrics from common result patterns 

662 if isinstance(tool_result, dict): 

663 # Look for common success indicators 

664 if tool_result.get("success") is True or "message" in tool_result: 

665 status_line += " [green]✓[/green]" 

666 elif tool_result.get("success") is False: 

667 status_line += " [red]✗[/red]" 

668 

669 # Extract key metrics based on common patterns 

670 if "total_count" in tool_result: 

671 metrics.append(f"{tool_result['total_count']} items") 

672 elif "count" in tool_result: 

673 metrics.append(f"{tool_result['count']} items") 

674 

675 # PII-specific metrics 

676 if "tables_with_pii" in tool_result: 

677 metrics.append(f"{tool_result['tables_with_pii']} tables with PII") 

678 if "total_pii_columns" in tool_result: 

679 metrics.append(f"{tool_result['total_pii_columns']} PII columns") 

680 

681 # Tag-specific metrics 

682 if "tagged_columns" in tool_result: 

683 metrics.append(f"{len(tool_result['tagged_columns'])} columns tagged") 

684 

685 # Status-specific info 

686 if tool_name == "status" and "workspace" in tool_result: 

687 workspace = tool_result.get("workspace", {}) 

688 if "name" in workspace: 

689 metrics.append(f"workspace: {workspace['name']}") 

690 

691 # Schema/Catalog selection specific info - keep it simple 

692 if tool_name == "set_schema" and "schema_name" in tool_result: 

693 metrics.append(f"{tool_result['schema_name']}") 

694 elif tool_name == "set_catalog" and "catalog_name" in tool_result: 

695 metrics.append(f"{tool_result['catalog_name']}") 

696 

697 # Generic message fallback 

698 if not metrics and "message" in tool_result: 

699 metrics.append(tool_result["message"]) 

700 

701 # Format the condensed display 

702 if metrics: 

703 status_line += f" ({', '.join(metrics)})" 

704 

705 self.console.print(status_line) 

706 

707 def _display_error(self, result: CommandResult) -> None: 

708 """Display an error from a command result.""" 

709 error_message = result.message or "Unknown error occurred" 

710 self.console.print(f"[{ERROR_STYLE}]Error: {error_message}[/{ERROR_STYLE}]") 

711 

712 if result.error and self.debug: 

713 # Use Rich's traceback rendering for better formatting 

714 self.console.print_exception(show_locals=True) 

715 

716 def _display_catalogs(self, data: Dict[str, Any]) -> None: 

717 """Display catalogs in a nicely formatted way.""" 

718 from src.ui.table_formatter import display_table 

719 from src.exceptions import PaginationCancelled 

720 

721 catalogs = data.get("catalogs", []) 

722 current_catalog = data.get("current_catalog") 

723 

724 if not catalogs: 

725 self.console.print(f"[{WARNING_STYLE}]No catalogs found.[/{WARNING_STYLE}]") 

726 # Raise PaginationCancelled to return to chuck > prompt immediately 

727 raise PaginationCancelled() 

728 

729 # Define column styling based on the active catalog 

730 def name_style(name): 

731 return "bold green" if name == current_catalog else None 

732 

733 style_map = {"name": name_style} 

734 

735 # Prepare catalog data - ensure lowercase types 

736 for catalog in catalogs: 

737 if "type" in catalog and catalog["type"]: 

738 catalog["type"] = catalog["type"].lower() 

739 

740 # Display the table 

741 display_table( 

742 console=self.console, 

743 data=catalogs, 

744 columns=["name", "type", "comment"], 

745 headers=["Name", "Type", "Comment"], 

746 title="Available Catalogs", 

747 style_map=style_map, 

748 title_style=TABLE_TITLE_STYLE, 

749 show_lines=False, 

750 ) 

751 

752 if current_catalog: 

753 self.console.print( 

754 f"\nCurrent catalog: [{SUCCESS_STYLE}]{current_catalog}[/{SUCCESS_STYLE}]" 

755 ) 

756 

757 # Raise PaginationCancelled to return to chuck > prompt immediately 

758 # This prevents agent from continuing processing after catalog display is complete 

759 raise PaginationCancelled() 

760 

761 def _display_schemas(self, data: Dict[str, Any]) -> None: 

762 """Display schemas in a nicely formatted way.""" 

763 from src.ui.table_formatter import display_table 

764 from src.exceptions import PaginationCancelled 

765 

766 schemas = data.get("schemas", []) 

767 catalog_name = data.get("catalog_name", "") 

768 current_schema = data.get("current_schema") 

769 

770 if not schemas: 

771 self.console.print( 

772 f"[{WARNING_STYLE}]No schemas found in catalog '{catalog_name}'.[/{WARNING_STYLE}]" 

773 ) 

774 # Raise PaginationCancelled to return to chuck > prompt immediately 

775 raise PaginationCancelled() 

776 

777 # Define column styling based on the active schema 

778 def name_style(name): 

779 return "bold green" if name == current_schema else None 

780 

781 style_map = {"name": name_style} 

782 

783 # Display the table 

784 display_table( 

785 console=self.console, 

786 data=schemas, 

787 columns=["name", "comment"], 

788 headers=["Name", "Comment"], 

789 title=f"Schemas in catalog '{catalog_name}'", 

790 style_map=style_map, 

791 title_style=TABLE_TITLE_STYLE, 

792 show_lines=False, 

793 ) 

794 

795 # Display current schema if available 

796 if current_schema: 

797 self.console.print( 

798 f"\nCurrent schema: [{SUCCESS_STYLE}]{current_schema}[/{SUCCESS_STYLE}]" 

799 ) 

800 

801 # Raise PaginationCancelled to return to chuck > prompt immediately 

802 # This prevents agent from continuing processing after schema display is complete 

803 raise PaginationCancelled() 

804 

805 def _display_tables(self, data: Dict[str, Any]) -> None: 

806 """Display tables in a nicely formatted way.""" 

807 from src.ui.table_formatter import display_table 

808 from src.exceptions import PaginationCancelled 

809 

810 tables = data.get("tables", []) 

811 catalog_name = data.get("catalog_name", "") 

812 schema_name = data.get("schema_name", "") 

813 total_count = data.get("total_count", len(tables)) 

814 

815 if not tables: 

816 self.console.print( 

817 f"[{WARNING_STYLE}]No tables found in {catalog_name}.{schema_name}[/{WARNING_STYLE}]" 

818 ) 

819 # Raise PaginationCancelled to return to chuck > prompt immediately 

820 raise PaginationCancelled() 

821 

822 # Process the table data for display 

823 for table in tables: 

824 # Convert columns list to count if present 

825 if "columns" in table and isinstance(table["columns"], list): 

826 table["column_count"] = len(table["columns"]) 

827 else: 

828 table["column_count"] = 0 

829 

830 # Format timestamps if present 

831 for ts_field in ["created_at", "updated_at"]: 

832 if ts_field in table and table[ts_field]: 

833 try: 

834 # Convert timestamp to more readable format if needed 

835 # Handle Unix timestamps (integers) and ISO strings 

836 timestamp = table[ts_field] 

837 if isinstance(timestamp, int): 

838 # Convert Unix timestamp (milliseconds) to readable date 

839 from datetime import datetime 

840 

841 date_obj = datetime.fromtimestamp(timestamp / 1000) 

842 table[ts_field] = date_obj.strftime("%Y-%m-%d") 

843 elif isinstance(timestamp, str) and len(timestamp) > 10: 

844 table[ts_field] = timestamp.split("T")[0] 

845 except Exception: 

846 pass # Keep the original format if conversion fails 

847 

848 # Format row count if present 

849 if "row_count" in table and table["row_count"] not in ["-", "Unknown"]: 

850 try: 

851 row_count = table["row_count"] 

852 if isinstance(row_count, str) and row_count.isdigit(): 

853 row_count = int(row_count) 

854 

855 if isinstance(row_count, int): 

856 # Format large numbers with appropriate suffixes 

857 if row_count >= 1_000_000_000: 

858 table["row_count"] = f"{row_count / 1_000_000_000:.1f}B" 

859 elif row_count >= 1_000_000: 

860 table["row_count"] = f"{row_count / 1_000_000:.1f}M" 

861 elif row_count >= 1_000: 

862 table["row_count"] = f"{row_count / 1_000:.1f}K" 

863 else: 

864 table["row_count"] = str(row_count) 

865 except Exception: 

866 pass # Keep the original format if conversion fails 

867 

868 # Define column styling functions 

869 def table_type_style(type_val): 

870 if type_val == "VIEW" or type_val == "view": 

871 return "bright_blue" 

872 return None 

873 

874 # Set up style map 

875 style_map = { 

876 "table_type": table_type_style, 

877 "column_count": lambda val: "dim" if val == 0 else None, 

878 } 

879 

880 # Adjust title based on method 

881 method = data.get("method", "") 

882 title = ( 

883 f"Tables in {catalog_name}.{schema_name} ({total_count} total)" 

884 if method == "unity_catalog" 

885 else "Available Tables" 

886 ) 

887 

888 # Display the table using our formatter 

889 display_table( 

890 console=self.console, 

891 data=tables, 

892 columns=[ 

893 "name", 

894 "table_type", 

895 "column_count", 

896 "row_count", 

897 "created_at", 

898 "updated_at", 

899 ], 

900 headers=["Table Name", "Type", "# Cols", "Rows", "Created", "Last Updated"], 

901 title=title, 

902 style_map=style_map, 

903 title_style=TABLE_TITLE_STYLE, 

904 show_lines=True, 

905 ) 

906 

907 # Raise PaginationCancelled to return to chuck > prompt immediately 

908 # This prevents agent from continuing processing after table display is complete 

909 raise PaginationCancelled() 

910 

911 def _display_models(self, models: List[Dict[str, Any]]) -> None: 

912 """Display models in a nicely formatted way.""" 

913 from src.ui.table_formatter import display_table 

914 from src.exceptions import PaginationCancelled 

915 

916 # Use imported function to get the active model for highlighting 

917 active_model = get_active_model() 

918 

919 if not models: 

920 self.console.print( 

921 f"[{WARNING_STYLE}]No models found or returned.[/{WARNING_STYLE}]" 

922 ) 

923 # Raise PaginationCancelled to return to chuck > prompt immediately 

924 raise PaginationCancelled() 

925 

926 # Process model data for display 

927 processed_models = [] 

928 for model in models: 

929 # Create a processed model with clear fields 

930 processed = { 

931 "name": model.get("name", "N/A"), 

932 "creator": model.get("creator", "N/A"), 

933 } 

934 

935 # Extract and process state information 

936 state = model.get("state", {}) 

937 ready_status = state.get("ready", "UNKNOWN").upper() 

938 processed["status"] = ready_status 

939 

940 # Add to our list 

941 processed_models.append(processed) 

942 

943 # Define styling function for status 

944 def status_style(status): 

945 if status == "READY": 

946 return "green" 

947 elif status == "NOT_READY": 

948 return "yellow" 

949 elif "ERROR" in status: 

950 return "red" 

951 return None 

952 

953 # Define styling function for the name to highlight active model 

954 def name_style(name): 

955 return "bold green" if name == active_model else None 

956 

957 # Process model names to add recommended tag 

958 for model in processed_models: 

959 if model["name"] in [ 

960 "databricks-meta-llama-3-3-70b-instruct", 

961 "databricks-claude-3-7-sonnet", 

962 ]: 

963 model["name"] = f"{model['name']} (recommended)" 

964 

965 # Set up style map 

966 style_map = {"name": name_style, "status": status_style} 

967 

968 # Display the table using our formatter 

969 display_table( 

970 console=self.console, 

971 data=processed_models, 

972 columns=["name", "creator", "status"], 

973 headers=["Endpoint Name", "Creator", "State"], 

974 title="Available Model Serving Endpoints", 

975 style_map=style_map, 

976 title_style=TABLE_TITLE_STYLE, 

977 show_lines=False, 

978 ) 

979 

980 # Display active model if set 

981 if active_model: 

982 self.console.print( 

983 f"\nCurrent active model: [{SUCCESS_STYLE}]{active_model}[/{SUCCESS_STYLE}]" 

984 ) 

985 

986 # Raise PaginationCancelled to return to chuck > prompt immediately 

987 # This prevents agent from continuing processing after model display is complete 

988 raise PaginationCancelled() 

989 

990 def _display_detailed_models(self, data: Dict[str, Any]) -> None: 

991 """Display models with detailed information and filtering.""" 

992 from src.ui.table_formatter import display_table 

993 from src.exceptions import PaginationCancelled 

994 

995 models = data.get("models", []) 

996 active_model = data.get("active_model") 

997 detailed = data.get("detailed", False) 

998 filter_text = data.get("filter") 

999 

1000 # If no models, display the help message 

1001 if not models: 

1002 self.console.print( 

1003 f"[{WARNING_STYLE}]No models found in workspace.[/{WARNING_STYLE}]" 

1004 ) 

1005 if data.get("message"): 

1006 self.console.print("\n" + data.get("message")) 

1007 # Raise PaginationCancelled to return to chuck > prompt immediately 

1008 raise PaginationCancelled() 

1009 

1010 # Display header with filter information if applicable 

1011 title = "Available Models" 

1012 if filter_text: 

1013 title += f" matching '{filter_text}'" 

1014 

1015 # Process model data for display 

1016 processed_models = [] 

1017 for model in models: 

1018 # Create a processed model entry 

1019 processed = { 

1020 "name": model.get("name", "N/A"), 

1021 "creator": model.get("creator", "N/A"), 

1022 } 

1023 

1024 # Get state information 

1025 state = model.get("state", {}) 

1026 ready_status = state.get("ready", "UNKNOWN").upper() 

1027 processed["status"] = ready_status 

1028 

1029 # Add detailed fields if requested 

1030 if detailed: 

1031 processed["endpoint_type"] = model.get("endpoint_type", "Unknown") 

1032 processed["last_modified"] = model.get("last_updated", "Unknown") 

1033 

1034 # Add any additional details from the details field 

1035 details = model.get("details", {}) 

1036 if details: 

1037 for key, value in details.items(): 

1038 # Only add if not already present and meaningful 

1039 if key not in processed and value is not None and key != "name": 

1040 processed[key] = value 

1041 

1042 # Add to our list 

1043 processed_models.append(processed) 

1044 

1045 # Define column styling functions 

1046 def status_style(status): 

1047 if status == "READY": 

1048 return "green" 

1049 elif status == "NOT_READY" or status == "UNKNOWN": 

1050 return "yellow" 

1051 elif "ERROR" in status: 

1052 return "red" 

1053 return None 

1054 

1055 # Define styling function for the name to highlight active model 

1056 def name_style(name): 

1057 if name == active_model: 

1058 return "bold green" 

1059 return "cyan" 

1060 

1061 # Set up style map with appropriate styles for each column 

1062 style_map = { 

1063 "name": name_style, 

1064 "status": status_style, 

1065 "creator": lambda _: f"{INFO}", 

1066 "endpoint_type": lambda _: f"{WARNING}", 

1067 "last_modified": lambda _: f"{DIALOG_BORDER}", 

1068 } 

1069 

1070 # Define columns and headers based on detail level 

1071 if detailed: 

1072 columns = ["name", "status", "creator", "endpoint_type", "last_modified"] 

1073 headers = ["Name", "Status", "Creator", "Type", "Last Modified"] 

1074 else: 

1075 columns = ["name", "status", "creator"] 

1076 headers = ["Model Name", "Status", "Creator"] 

1077 

1078 # Display the table using our formatter 

1079 display_table( 

1080 console=self.console, 

1081 data=processed_models, 

1082 columns=columns, 

1083 headers=headers, 

1084 title=title, 

1085 style_map=style_map, 

1086 title_style=TABLE_TITLE_STYLE, 

1087 show_lines=False, 

1088 box_style="SIMPLE", 

1089 ) 

1090 

1091 # Display help instructions 

1092 self.console.print( 

1093 "\n[dim]Use [bold]/select_model <name>[/bold] to set the active model for agent operations.[/dim]" 

1094 ) 

1095 if not detailed: 

1096 self.console.print( 

1097 "[dim]Use [bold]/list-models --detailed[/bold] to see more details about available models.[/dim]" 

1098 ) 

1099 if not filter_text: 

1100 self.console.print( 

1101 "[dim]Use [bold]/list-models --filter <text>[/bold] to filter models by name.[/dim]" 

1102 ) 

1103 self.console.print( 

1104 "[dim]Use [bold]/models[/bold] for a simpler view of available models.[/dim]" 

1105 ) 

1106 

1107 # Raise PaginationCancelled to return to chuck > prompt immediately 

1108 # This prevents agent from continuing processing after detailed model display is complete 

1109 raise PaginationCancelled() 

1110 

1111 def _display_warehouses(self, data: Dict[str, Any]) -> None: 

1112 """Display SQL warehouses in a nicely formatted way.""" 

1113 from src.ui.table_formatter import display_table 

1114 from src.exceptions import PaginationCancelled 

1115 

1116 warehouses = data.get("warehouses", []) 

1117 current_warehouse_id = data.get("current_warehouse_id") 

1118 

1119 if not warehouses: 

1120 self.console.print( 

1121 f"[{WARNING_STYLE}]No SQL warehouses found.[/{WARNING_STYLE}]" 

1122 ) 

1123 # Raise PaginationCancelled to return to chuck > prompt immediately 

1124 raise PaginationCancelled() 

1125 

1126 # Process warehouse data for display 

1127 processed_warehouses = [] 

1128 for warehouse in warehouses: 

1129 # Create a processed warehouse with formatted fields 

1130 processed = { 

1131 "name": warehouse.get("name", ""), 

1132 "id": warehouse.get("id", ""), 

1133 "size": warehouse.get("cluster_size", ""), # API uses cluster_size 

1134 "state": warehouse.get("state", ""), 

1135 } 

1136 processed_warehouses.append(processed) 

1137 

1138 # Define styling function for name based on current warehouse 

1139 def name_style(name, row): 

1140 if row.get("id") == current_warehouse_id: 

1141 return "bold green" 

1142 return None 

1143 

1144 # Define styling function for ID based on current warehouse 

1145 def id_style(id_val): 

1146 if id_val == current_warehouse_id: 

1147 return "bold green" 

1148 return None 

1149 

1150 # Define styling function for state 

1151 def state_style(state): 

1152 if state == "RUNNING": 

1153 return "green" 

1154 elif state == "STOPPED": 

1155 return "red" 

1156 elif state in ["STARTING", "STOPPING", "DELETING", "RESIZING"]: 

1157 return "yellow" 

1158 return "dim" 

1159 

1160 # Set up style map 

1161 style_map = { 

1162 "name": lambda name, row=None: name_style(name, row), 

1163 "id": id_style, 

1164 "state": state_style, 

1165 } 

1166 

1167 # Set maximum lengths for fields 

1168 

1169 # Display the table 

1170 display_table( 

1171 console=self.console, 

1172 data=processed_warehouses, 

1173 columns=["name", "id", "size", "state"], 

1174 headers=["Name", "ID", "Size", "State"], 

1175 title="Available SQL Warehouses", 

1176 style_map=style_map, 

1177 title_style=TABLE_TITLE_STYLE, 

1178 show_lines=False, 

1179 ) 

1180 

1181 # Display current warehouse ID if set 

1182 if current_warehouse_id: 

1183 self.console.print( 

1184 f"\nCurrent SQL warehouse ID: [{SUCCESS_STYLE}]{current_warehouse_id}[/{SUCCESS_STYLE}]" 

1185 ) 

1186 

1187 # Raise PaginationCancelled to return to chuck > prompt immediately 

1188 # This prevents agent from continuing processing after warehouse display is complete 

1189 raise PaginationCancelled() 

1190 

1191 def _display_volumes(self, data: Dict[str, Any]) -> None: 

1192 """Display volumes in a nicely formatted way.""" 

1193 from src.ui.table_formatter import display_table 

1194 from src.exceptions import PaginationCancelled 

1195 

1196 volumes = data.get("volumes", []) 

1197 catalog_name = data.get("catalog_name", "") 

1198 schema_name = data.get("schema_name", "") 

1199 

1200 if not volumes: 

1201 self.console.print( 

1202 f"[{WARNING_STYLE}]No volumes found in {catalog_name}.{schema_name}.[/{WARNING_STYLE}]" 

1203 ) 

1204 # Raise PaginationCancelled to return to chuck > prompt immediately 

1205 raise PaginationCancelled() 

1206 

1207 # Process volume data for display 

1208 processed_volumes = [] 

1209 for volume in volumes: 

1210 # Create a processed volume with normalized fields 

1211 processed = { 

1212 "name": volume.get("name", ""), 

1213 "type": volume.get( 

1214 "volume_type", "" 

1215 ).upper(), # Use upper for consistency 

1216 "comment": volume.get("comment", ""), 

1217 } 

1218 processed_volumes.append(processed) 

1219 

1220 # Define styling for volume types 

1221 def type_style(volume_type): 

1222 if volume_type == "EXTERNAL": # Example conditional styling 

1223 return "yellow" 

1224 elif volume_type == "MANAGED": # Example conditional styling 

1225 return "blue" 

1226 return None 

1227 

1228 # Set up style map 

1229 style_map = {"type": type_style} 

1230 

1231 # Display the table 

1232 display_table( 

1233 console=self.console, 

1234 data=processed_volumes, 

1235 columns=["name", "type", "comment"], 

1236 headers=["Name", "Type", "Comment"], 

1237 title=f"Volumes in {catalog_name}.{schema_name}", 

1238 style_map=style_map, 

1239 title_style=TABLE_TITLE_STYLE, 

1240 show_lines=False, 

1241 ) 

1242 

1243 # Raise PaginationCancelled to return to chuck > prompt immediately 

1244 # This prevents agent from continuing processing after volume display is complete 

1245 raise PaginationCancelled() 

1246 

1247 def _display_status(self, data: Dict[str, Any]) -> None: 

1248 """Display current status information including connection status and permissions.""" 

1249 from src.ui.table_formatter import display_table 

1250 

1251 workspace_url = data.get("workspace_url", "Not set") 

1252 active_catalog = data.get("active_catalog", "Not set") 

1253 active_schema = data.get("active_schema", "Not set") 

1254 active_model = data.get("active_model", "Not set") 

1255 warehouse_id = data.get("warehouse_id", "Not set") 

1256 connection_status = data.get("connection_status", "Unknown") 

1257 

1258 # Prepare settings for display 

1259 status_items = [ 

1260 {"setting": "Workspace URL", "value": workspace_url}, 

1261 {"setting": "Active Catalog", "value": active_catalog}, 

1262 {"setting": "Active Schema", "value": active_schema}, 

1263 {"setting": "Active Model", "value": active_model}, 

1264 {"setting": "Active Warehouse", "value": warehouse_id}, 

1265 {"setting": "Connection Status", "value": connection_status}, 

1266 ] 

1267 

1268 # Define styling functions 

1269 def value_style(value, row): 

1270 setting = row.get("setting", "") 

1271 

1272 # Special handling for connection status 

1273 if setting == "Connection Status": 

1274 if value == "Connected - token is valid": 

1275 return "green" 

1276 elif "Invalid" in value or "Not connected" in value: 

1277 return "red" 

1278 else: 

1279 return "yellow" 

1280 # General styling for values 

1281 elif value != "Not set": 

1282 return "green" 

1283 else: 

1284 return "yellow" 

1285 

1286 # Set up style map 

1287 style_map = {"value": lambda value, row: value_style(value, row)} 

1288 

1289 # Display the status table 

1290 display_table( 

1291 console=self.console, 

1292 data=status_items, 

1293 columns=["setting", "value"], 

1294 headers=["Setting", "Value"], 

1295 title="Current Configuration", 

1296 style_map=style_map, 

1297 title_style=TABLE_TITLE_STYLE, 

1298 show_lines=False, 

1299 ) 

1300 

1301 # If permissions data is available, display it 

1302 permissions_data = data.get("permissions") 

1303 if permissions_data: 

1304 self._display_permissions(permissions_data) 

1305 

1306 def _display_permissions(self, permissions_data: Dict[str, Any]) -> None: 

1307 """ 

1308 Display detailed permission check results. 

1309 

1310 Args: 

1311 permissions_data: Dictionary of permission check results 

1312 """ 

1313 from src.ui.table_formatter import display_table 

1314 

1315 if not permissions_data: 

1316 self.console.print( 

1317 f"[{WARNING_STYLE}]No permission data available.[/{WARNING_STYLE}]" 

1318 ) 

1319 return 

1320 

1321 # Format permission data for display 

1322 formatted_permissions = [] 

1323 for resource, data in permissions_data.items(): 

1324 authorized = data.get("authorized", False) 

1325 details = ( 

1326 data.get("details") 

1327 if authorized 

1328 else data.get("error", "Access denied") 

1329 ) 

1330 api_path = data.get("api_path", "Unknown") 

1331 

1332 # Create a dictionary for this permission 

1333 resource_name = resource.replace("_", " ").title() 

1334 permission_entry = { 

1335 "resource": resource_name, 

1336 "status": "Authorized" if authorized else "Denied", 

1337 "details": details, 

1338 "api_path": api_path, # Store for reference in the endpoints section 

1339 "authorized": authorized, # Store for conditional styling 

1340 } 

1341 formatted_permissions.append(permission_entry) 

1342 

1343 # Define styling function for status column 

1344 def status_style(status, row): 

1345 return "green" if row.get("authorized") else "red" 

1346 

1347 # Set up style map 

1348 style_map = {"status": status_style} 

1349 

1350 # Display the permissions table 

1351 display_table( 

1352 console=self.console, 

1353 data=formatted_permissions, 

1354 columns=["resource", "status", "details"], 

1355 headers=["Resource", "Status", "Details"], 

1356 title="Databricks API Token Permissions", 

1357 style_map=style_map, 

1358 title_style=TABLE_TITLE_STYLE, 

1359 show_lines=True, 

1360 ) 

1361 

1362 # Additional note about API endpoints 

1363 self.console.print("\n[dim]API endpoints checked:[/dim]") 

1364 for item in formatted_permissions: 

1365 resource_name = item["resource"] 

1366 api_path = item["api_path"] 

1367 self.console.print(f"[dim]- {resource_name}: {api_path}[/dim]") 

1368 

1369 def _display_table_details(self, data: Dict[str, Any]) -> None: 

1370 """Display detailed information for a single table.""" 

1371 from src.ui.table_formatter import display_table 

1372 

1373 table = data.get("table", {}) 

1374 full_name = data.get("full_name", "") 

1375 has_delta_metadata = data.get("has_delta_metadata", False) 

1376 

1377 if not table: 

1378 self.console.print( 

1379 f"[{WARNING_STYLE}]No table details available.[/{WARNING_STYLE}]" 

1380 ) 

1381 return 

1382 

1383 # Display table header 

1384 self.console.print( 

1385 f"\n[{TABLE_TITLE_STYLE}]Table Details: {full_name}[/{TABLE_TITLE_STYLE}]" 

1386 ) 

1387 

1388 # Prepare basic information data 

1389 basic_info = [] 

1390 properties = [ 

1391 ("Name", table.get("name", "")), 

1392 ("Full Name", full_name), 

1393 ("Type", table.get("table_type", "")), 

1394 ("Format", table.get("data_source_format", "")), 

1395 ("Storage Location", table.get("storage_location", "")), 

1396 ("Owner", table.get("owner", "")), 

1397 ("Created", table.get("created_at", "")), 

1398 ("Created By", table.get("created_by", "")), 

1399 ("Updated", table.get("updated_at", "")), 

1400 ("Updated By", table.get("updated_by", "")), 

1401 ("Comment", table.get("comment", "")), 

1402 ] 

1403 

1404 for prop, value in properties: 

1405 if value: # Only include non-empty values 

1406 basic_info.append({"property": prop, "value": value}) 

1407 

1408 # Display basic information table 

1409 self.console.print("\n[bold]Basic Information:[/bold]") 

1410 display_table( 

1411 console=self.console, 

1412 data=basic_info, 

1413 columns=["property", "value"], 

1414 headers=["Property", "Value"], 

1415 show_lines=False, 

1416 ) 

1417 

1418 # Display columns if available 

1419 columns_data = table.get("columns", []) 

1420 if columns_data: 

1421 # Prepare column data 

1422 columns_for_display = [] 

1423 for column in columns_data: 

1424 columns_for_display.append( 

1425 { 

1426 "name": column.get("name", ""), 

1427 "type": column.get("type_text", column.get("type", "")), 

1428 "nullable": "Yes" if column.get("nullable", False) else "No", 

1429 "comment": column.get("comment", ""), 

1430 } 

1431 ) 

1432 

1433 # Display columns table 

1434 self.console.print("\n[bold]Columns:[/bold]") 

1435 display_table( 

1436 console=self.console, 

1437 data=columns_for_display, 

1438 columns=["name", "type", "nullable", "comment"], 

1439 headers=["Name", "Type", "Nullable", "Comment"], 

1440 show_lines=False, 

1441 ) 

1442 

1443 # Display properties if available 

1444 properties_data = table.get("properties", {}) 

1445 if properties_data: 

1446 # Prepare properties data 

1447 props_for_display = [] 

1448 for prop, value in properties_data.items(): 

1449 # Skip empty values 

1450 if value is None or value == "": 

1451 continue 

1452 

1453 props_for_display.append({"property": prop, "value": value}) 

1454 

1455 # Display properties table 

1456 self.console.print("\n[bold]Table Properties:[/bold]") 

1457 display_table( 

1458 console=self.console, 

1459 data=props_for_display, 

1460 columns=["property", "value"], 

1461 headers=["Property", "Value"], 

1462 show_lines=False, 

1463 ) 

1464 

1465 # Display Delta metadata if available 

1466 if has_delta_metadata and "delta" in table: 

1467 delta_info = table.get("delta", {}) 

1468 

1469 # Prepare Delta metadata data 

1470 delta_for_display = [] 

1471 delta_properties = [ 

1472 ("Format", delta_info.get("format", "")), 

1473 ("ID", delta_info.get("id", "")), 

1474 ("Last Updated", delta_info.get("last_updated", "")), 

1475 ("Min Reader Version", delta_info.get("min_reader_version", "")), 

1476 ("Min Writer Version", delta_info.get("min_writer_version", "")), 

1477 ("Num Files", delta_info.get("num_files", "")), 

1478 ("Size (Bytes)", delta_info.get("size_in_bytes", "")), 

1479 ] 

1480 

1481 for prop, value in delta_properties: 

1482 if value: # Only include non-empty values 

1483 delta_for_display.append({"property": prop, "value": value}) 

1484 

1485 # Display Delta metadata table 

1486 if delta_for_display: # Only if we have data to show 

1487 self.console.print("\n[bold]Delta Metadata:[/bold]") 

1488 display_table( 

1489 console=self.console, 

1490 data=delta_for_display, 

1491 columns=["property", "value"], 

1492 headers=["Property", "Value"], 

1493 show_lines=False, 

1494 ) 

1495 

1496 def _display_catalog_details(self, data: Dict[str, Any]) -> None: 

1497 """Display detailed information for a specific catalog.""" 

1498 from src.ui.table_formatter import display_table 

1499 

1500 catalog = data 

1501 

1502 if not catalog: 

1503 self.console.print( 

1504 f"[{WARNING_STYLE}]No catalog details available.[/{WARNING_STYLE}]" 

1505 ) 

1506 return 

1507 

1508 # Display catalog header 

1509 catalog_name = catalog.get("name", "Unknown") 

1510 self.console.print( 

1511 f"\n[{TABLE_TITLE_STYLE}]Catalog Details: {catalog_name}[/{TABLE_TITLE_STYLE}]" 

1512 ) 

1513 

1514 # Prepare basic information data 

1515 basic_info = [] 

1516 properties = [ 

1517 ("Name", catalog.get("name", "")), 

1518 ("Type", catalog.get("type", "")), 

1519 ("Comment", catalog.get("comment", "")), 

1520 ("Provider", catalog.get("provider", {}).get("name", "")), 

1521 ("Storage Root", catalog.get("storage_root", "")), 

1522 ("Storage Location", catalog.get("storage_location", "")), 

1523 ("Owner", catalog.get("owner", "")), 

1524 ("Created At", catalog.get("created_at", "")), 

1525 ("Created By", catalog.get("created_by", "")), 

1526 ("Options", str(catalog.get("options", {}))), 

1527 ] 

1528 

1529 for prop, value in properties: 

1530 if value: # Only include non-empty values 

1531 basic_info.append({"property": prop, "value": value}) 

1532 

1533 # Display basic information table 

1534 display_table( 

1535 console=self.console, 

1536 data=basic_info, 

1537 columns=["property", "value"], 

1538 headers=["Property", "Value"], 

1539 show_lines=False, 

1540 ) 

1541 

1542 def _display_schema_details(self, data: Dict[str, Any]) -> None: 

1543 """Display detailed information for a specific schema.""" 

1544 from src.ui.table_formatter import display_table 

1545 

1546 schema = data 

1547 

1548 if not schema: 

1549 self.console.print( 

1550 f"[{WARNING_STYLE}]No schema details available.[/{WARNING_STYLE}]" 

1551 ) 

1552 return 

1553 

1554 # Display schema header 

1555 schema_name = schema.get("name", "Unknown") 

1556 catalog_name = schema.get("catalog_name", "Unknown") 

1557 full_name = f"{catalog_name}.{schema_name}" 

1558 self.console.print( 

1559 f"\n[{TABLE_TITLE_STYLE}]Schema Details: {full_name}[/{TABLE_TITLE_STYLE}]" 

1560 ) 

1561 

1562 # Prepare basic information data 

1563 basic_info = [] 

1564 properties = [ 

1565 ("Name", schema.get("name", "")), 

1566 ("Full Name", schema.get("full_name", "")), 

1567 ("Catalog Name", schema.get("catalog_name", "")), 

1568 ("Comment", schema.get("comment", "")), 

1569 ("Storage Root", schema.get("storage_root", "")), 

1570 ("Storage Location", schema.get("storage_location", "")), 

1571 ("Owner", schema.get("owner", "")), 

1572 ("Created At", schema.get("created_at", "")), 

1573 ("Created By", schema.get("created_by", "")), 

1574 ] 

1575 

1576 for prop, value in properties: 

1577 if value: # Only include non-empty values 

1578 basic_info.append({"property": prop, "value": value}) 

1579 

1580 # Display basic information table 

1581 display_table( 

1582 console=self.console, 

1583 data=basic_info, 

1584 columns=["property", "value"], 

1585 headers=["Property", "Value"], 

1586 show_lines=False, 

1587 ) 

1588 

1589 def _display_pii_scan_results(self, data: Dict[str, Any]) -> None: 

1590 """Display PII scan results for tables in a schema.""" 

1591 from src.ui.table_formatter import display_table 

1592 

1593 if not data: 

1594 self.console.print( 

1595 f"[{WARNING_STYLE}]No PII scan results available.[/{WARNING_STYLE}]" 

1596 ) 

1597 return 

1598 

1599 catalog_name = data.get("catalog", "Unknown") 

1600 schema_name = data.get("schema", "Unknown") 

1601 results_detail = data.get("results_detail", []) 

1602 tables_with_pii = data.get("tables_with_pii", 0) 

1603 total_pii_columns = data.get("total_pii_columns", 0) 

1604 

1605 # Display summary header 

1606 self.console.print( 

1607 f"\n[{TABLE_TITLE_STYLE}]PII Scan Results: {catalog_name}.{schema_name}[/{TABLE_TITLE_STYLE}]" 

1608 ) 

1609 self.console.print( 

1610 f"Found {tables_with_pii} tables with a total of {total_pii_columns} PII columns." 

1611 ) 

1612 

1613 # If no details, exit early 

1614 if not results_detail: 

1615 self.console.print( 

1616 f"[{WARNING_STYLE}]No detailed scan results available.[/{WARNING_STYLE}]" 

1617 ) 

1618 return 

1619 

1620 # Prepare table data for display - just tables with PII 

1621 tables_with_pii_data = [] 

1622 for table_result in results_detail: 

1623 if not table_result.get("skipped", False) and table_result.get( 

1624 "has_pii", False 

1625 ): 

1626 # Format data for display 

1627 table_data = { 

1628 "name": table_result.get("table_name", ""), 

1629 "full_name": table_result.get("full_name", ""), 

1630 "pii_columns_count": table_result.get("pii_column_count", 0), 

1631 "total_columns": table_result.get("column_count", 0), 

1632 } 

1633 tables_with_pii_data.append(table_data) 

1634 

1635 # Sort by PII column count, most PII columns first 

1636 tables_with_pii_data.sort( 

1637 key=lambda x: x.get("pii_columns_count", 0), reverse=True 

1638 ) 

1639 

1640 # Display tables with PII 

1641 if tables_with_pii_data: 

1642 self.console.print("\n[bold]Tables with PII:[/bold]") 

1643 display_table( 

1644 console=self.console, 

1645 data=tables_with_pii_data, 

1646 columns=["name", "full_name", "pii_columns_count", "total_columns"], 

1647 headers=["Table Name", "Full Name", "PII Columns", "Total Columns"], 

1648 title="Tables with PII", 

1649 title_style=TABLE_TITLE_STYLE, 

1650 show_lines=True, 

1651 ) 

1652 

1653 # For each table with PII, display the PII columns 

1654 for table_result in results_detail: 

1655 if not table_result.get("skipped", False) and table_result.get( 

1656 "has_pii", False 

1657 ): 

1658 table_name = table_result.get("table_name", "") 

1659 pii_columns = table_result.get("pii_columns", []) 

1660 

1661 if pii_columns: 

1662 self.console.print( 

1663 f"\n[bold]PII Columns in {table_name}:[/bold]" 

1664 ) 

1665 

1666 # Prepare column data 

1667 column_data = [] 

1668 for col in pii_columns: 

1669 column_data.append( 

1670 { 

1671 "name": col.get("name", ""), 

1672 "type": col.get("type", ""), 

1673 "semantic": col.get("semantic", ""), 

1674 } 

1675 ) 

1676 

1677 # Display column data 

1678 display_table( 

1679 console=self.console, 

1680 data=column_data, 

1681 columns=["name", "type", "semantic"], 

1682 headers=["Column Name", "Data Type", "PII Type"], 

1683 show_lines=False, 

1684 ) 

1685 else: 

1686 self.console.print( 

1687 f"[{WARNING_STYLE}]No tables with PII columns found.[/{WARNING_STYLE}]" 

1688 ) 

1689 

1690 def _display_sql_results(self, data: Dict[str, Any]) -> None: 

1691 """Display SQL query results in a formatted table.""" 

1692 from src.ui.table_formatter import display_table 

1693 from src.exceptions import PaginationCancelled 

1694 

1695 if not data: 

1696 self.console.print( 

1697 f"[{WARNING_STYLE}]No SQL results available.[/{WARNING_STYLE}]" 

1698 ) 

1699 return 

1700 

1701 columns = data.get("columns", []) 

1702 rows = data.get("rows", []) 

1703 row_count = data.get("row_count", 0) 

1704 execution_time = data.get("execution_time_ms") 

1705 

1706 if not rows: 

1707 self.console.print( 

1708 f"[{WARNING_STYLE}]Query returned no results.[/{WARNING_STYLE}]" 

1709 ) 

1710 return 

1711 

1712 # Check if we should paginate - either external links OR > 50 rows 

1713 should_paginate = ( 

1714 data.get("is_paginated", False) # External links case 

1715 or len(rows) > 50 # Large result set in data_array 

1716 ) 

1717 

1718 if should_paginate: 

1719 self._display_paginated_sql_results_local(data) 

1720 return 

1721 

1722 # Small result set - display normally 

1723 # Convert rows (list of lists) to list of dictionaries for display_table 

1724 formatted_data = [] 

1725 for row in rows: 

1726 row_dict = {} 

1727 for i, value in enumerate(row): 

1728 if i < len(columns): 

1729 row_dict[columns[i]] = value if value is not None else "" 

1730 formatted_data.append(row_dict) 

1731 

1732 # Create title with execution info 

1733 title = f"SQL Query Results ({row_count} rows" 

1734 if execution_time is not None: 

1735 title += f", {execution_time}ms" 

1736 title += ")" 

1737 

1738 # Display the results table 

1739 display_table( 

1740 console=self.console, 

1741 data=formatted_data, 

1742 columns=columns, 

1743 headers=columns, 

1744 title=title, 

1745 title_style=TABLE_TITLE_STYLE, 

1746 show_lines=True, 

1747 ) 

1748 

1749 # Raise PaginationCancelled to return to chuck > prompt immediately 

1750 # This prevents agent from continuing processing after SQL display is complete 

1751 raise PaginationCancelled() 

1752 

1753 def _display_sql_results_formatted(self, data: Dict[str, Any]) -> None: 

1754 """Display SQL query results from the original command result data.""" 

1755 # Since we now pass original data to TUI, we can use the regular display method 

1756 self._display_sql_results(data) 

1757 

1758 def _display_paginated_sql_results(self, data: Dict[str, Any]) -> None: 

1759 """Display paginated SQL query results with interactive navigation.""" 

1760 import sys 

1761 from src.commands.sql_external_data import PaginatedSQLResult 

1762 from src.ui.table_formatter import display_table 

1763 from src.exceptions import PaginationCancelled 

1764 

1765 columns = data.get("columns", []) 

1766 external_links = data.get("external_links", []) 

1767 total_row_count = data.get("total_row_count", 0) 

1768 chunks = data.get("chunks", []) 

1769 execution_time = data.get("execution_time_ms") 

1770 

1771 if not external_links: 

1772 self.console.print( 

1773 f"[{WARNING_STYLE}]No external data links available.[/{WARNING_STYLE}]" 

1774 ) 

1775 return 

1776 

1777 # Initialize paginated result handler 

1778 paginated_result = PaginatedSQLResult( 

1779 columns=columns, 

1780 external_links=external_links, 

1781 total_row_count=total_row_count, 

1782 chunks=chunks, 

1783 ) 

1784 

1785 rows_displayed = 0 

1786 page_num = 1 

1787 

1788 try: 

1789 while True: 

1790 # Get the next page of data 

1791 try: 

1792 rows, has_more = paginated_result.get_next_page() 

1793 except Exception as e: 

1794 self.console.print( 

1795 f"[{ERROR_STYLE}]Error fetching data: {str(e)}[/{ERROR_STYLE}]" 

1796 ) 

1797 break 

1798 

1799 if not rows and rows_displayed == 0: 

1800 self.console.print( 

1801 f"[{WARNING_STYLE}]Query returned no results.[/{WARNING_STYLE}]" 

1802 ) 

1803 break 

1804 

1805 if rows: 

1806 # Convert rows to the format expected by display_table 

1807 formatted_data = [] 

1808 for row in rows: 

1809 row_dict = {} 

1810 for i, value in enumerate(row): 

1811 if i < len(columns): 

1812 row_dict[columns[i]] = ( 

1813 value if value is not None else "" 

1814 ) 

1815 formatted_data.append(row_dict) 

1816 

1817 # Create title with pagination info 

1818 start_row = rows_displayed + 1 

1819 end_row = rows_displayed + len(rows) 

1820 title = f"SQL Query Results (Rows {start_row}-{end_row} of {total_row_count}" 

1821 if execution_time is not None and page_num == 1: 

1822 title += f", {execution_time}ms" 

1823 title += ")" 

1824 

1825 # Display the current page 

1826 display_table( 

1827 console=self.console, 

1828 data=formatted_data, 

1829 columns=columns, 

1830 headers=columns, 

1831 title=title, 

1832 title_style=TABLE_TITLE_STYLE, 

1833 show_lines=True, 

1834 ) 

1835 

1836 rows_displayed += len(rows) 

1837 page_num += 1 

1838 

1839 # Check if there are more pages 

1840 if not has_more: 

1841 self.console.print( 

1842 f"\n[{INFO_STYLE}]End of results ({total_row_count} total rows)[/{INFO_STYLE}]" 

1843 ) 

1844 # Raise PaginationCancelled to return to chuck > prompt immediately 

1845 raise PaginationCancelled() 

1846 

1847 # Show pagination prompt 

1848 self.console.print( 

1849 f"\n[dim]Press [bold]SPACE[/bold] for next page, or [bold]q[/bold] to quit... ({rows_displayed}/{total_row_count} rows shown)[/dim]" 

1850 ) 

1851 

1852 # Get user input 

1853 try: 

1854 if sys.stdin.isatty(): 

1855 import readchar 

1856 

1857 char = readchar.readchar() 

1858 

1859 # Handle user input 

1860 if char.lower() == "q": 

1861 raise PaginationCancelled() 

1862 elif char == " ": 

1863 # Continue to next page 

1864 self.console.print() # Add spacing 

1865 continue 

1866 else: 

1867 # Invalid input, show help 

1868 self.console.print( 

1869 f"[{WARNING_STYLE}]Press SPACE for next page or 'q' to quit[/{WARNING_STYLE}]" 

1870 ) 

1871 continue 

1872 else: 

1873 # Not a TTY (e.g., running in a script), auto-continue 

1874 self.console.print( 

1875 "[dim]Auto-continuing (not in interactive terminal)...[/dim]" 

1876 ) 

1877 continue 

1878 

1879 except (KeyboardInterrupt, EOFError): 

1880 raise # Re-raise to bubble up to main TUI loop 

1881 except Exception as e: 

1882 if isinstance(e, PaginationCancelled): 

1883 raise # Re-raise PaginationCancelled to bubble up 

1884 self.console.print( 

1885 f"\n[{ERROR_STYLE}]Input error: {str(e)}[/{ERROR_STYLE}]" 

1886 ) 

1887 # Fall back to regular input 

1888 try: 

1889 response = ( 

1890 input( 

1891 "[dim]Type 'q' to quit or press ENTER to continue: [/dim]" 

1892 ) 

1893 .strip() 

1894 .lower() 

1895 ) 

1896 if response == "q": 

1897 raise PaginationCancelled() 

1898 else: 

1899 continue 

1900 except (KeyboardInterrupt, EOFError): 

1901 raise # Re-raise to bubble up to main TUI loop 

1902 

1903 except Exception as e: 

1904 if isinstance(e, PaginationCancelled): 

1905 raise # Re-raise PaginationCancelled to bubble up 

1906 self.console.print( 

1907 f"[{ERROR_STYLE}]Error during pagination: {str(e)}[/{ERROR_STYLE}]" 

1908 ) 

1909 

1910 def _display_paginated_sql_results_local(self, data: Dict[str, Any]) -> None: 

1911 """Display paginated SQL query results with interactive navigation for local data.""" 

1912 import sys 

1913 from src.ui.table_formatter import display_table 

1914 from src.exceptions import PaginationCancelled 

1915 

1916 columns = data.get("columns", []) 

1917 rows = data.get("rows", []) 

1918 execution_time = data.get("execution_time_ms") 

1919 total_rows = len(rows) 

1920 

1921 # Check if this has external links (true pagination) or local rows (chunked display) 

1922 if data.get("is_paginated", False) and data.get("external_links"): 

1923 # Use the existing external links pagination 

1924 self._display_paginated_sql_results(data) 

1925 return 

1926 

1927 # Local pagination for large row sets 

1928 page_size = 50 

1929 current_position = 0 

1930 page_num = 1 

1931 

1932 try: 

1933 while current_position < total_rows: 

1934 # Get current page of rows 

1935 end_position = min(current_position + page_size, total_rows) 

1936 page_rows = rows[current_position:end_position] 

1937 

1938 if page_rows: 

1939 # Convert rows to the format expected by display_table 

1940 formatted_data = [] 

1941 for row in page_rows: 

1942 row_dict = {} 

1943 for i, value in enumerate(row): 

1944 if i < len(columns): 

1945 row_dict[columns[i]] = ( 

1946 value if value is not None else "" 

1947 ) 

1948 formatted_data.append(row_dict) 

1949 

1950 # Create title with pagination info 

1951 start_row = current_position + 1 

1952 title = f"SQL Query Results (Rows {start_row}-{end_position} of {total_rows}" 

1953 if execution_time is not None and page_num == 1: 

1954 title += f", {execution_time}ms" 

1955 title += ")" 

1956 

1957 # Display the current page 

1958 display_table( 

1959 console=self.console, 

1960 data=formatted_data, 

1961 columns=columns, 

1962 headers=columns, 

1963 title=title, 

1964 title_style=TABLE_TITLE_STYLE, 

1965 show_lines=True, 

1966 ) 

1967 

1968 current_position = end_position 

1969 page_num += 1 

1970 

1971 # Check if there are more pages 

1972 if current_position >= total_rows: 

1973 self.console.print( 

1974 f"\n[{INFO_STYLE}]End of results ({total_rows} total rows)[/{INFO_STYLE}]" 

1975 ) 

1976 # Raise PaginationCancelled to return to chuck > prompt immediately 

1977 raise PaginationCancelled() 

1978 

1979 # Show pagination prompt 

1980 self.console.print( 

1981 f"\n[dim]Press [bold]SPACE[/bold] for next page, or [bold]q[/bold] to quit... ({current_position}/{total_rows} rows shown)[/dim]" 

1982 ) 

1983 

1984 # Get user input 

1985 try: 

1986 if sys.stdin.isatty(): 

1987 import readchar 

1988 

1989 char = readchar.readchar() 

1990 

1991 # Handle user input 

1992 if char.lower() == "q": 

1993 raise PaginationCancelled() 

1994 elif char == " ": 

1995 # Continue to next page 

1996 self.console.print() # Add spacing 

1997 continue 

1998 else: 

1999 # Invalid input, show help 

2000 self.console.print( 

2001 f"[{WARNING_STYLE}]Press SPACE for next page or 'q' to quit[/{WARNING_STYLE}]" 

2002 ) 

2003 continue 

2004 else: 

2005 # Not a TTY (e.g., running in a script), auto-continue 

2006 self.console.print( 

2007 "[dim]Auto-continuing (not in interactive terminal)...[/dim]" 

2008 ) 

2009 continue 

2010 

2011 except (KeyboardInterrupt, EOFError): 

2012 raise # Re-raise to bubble up to main TUI loop 

2013 except Exception as e: 

2014 if isinstance(e, PaginationCancelled): 

2015 raise # Re-raise PaginationCancelled to bubble up 

2016 self.console.print( 

2017 f"\n[{ERROR_STYLE}]Input error: {str(e)}[/{ERROR_STYLE}]" 

2018 ) 

2019 # Fall back to regular input 

2020 try: 

2021 response = ( 

2022 input( 

2023 "[dim]Type 'q' to quit or press ENTER to continue: [/dim]" 

2024 ) 

2025 .strip() 

2026 .lower() 

2027 ) 

2028 if response == "q": 

2029 raise PaginationCancelled() 

2030 else: 

2031 continue 

2032 except (KeyboardInterrupt, EOFError): 

2033 raise # Re-raise to bubble up to main TUI loop 

2034 

2035 except Exception as e: 

2036 if isinstance(e, PaginationCancelled): 

2037 raise # Re-raise PaginationCancelled to bubble up 

2038 self.console.print( 

2039 f"[{ERROR_STYLE}]Error during pagination: {str(e)}[/{ERROR_STYLE}]" 

2040 )