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
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 23:16 -0700
1"""
2Main TUI interface for CHUCK AI.
3"""
5import os
6import shlex
7from typing import List, Dict, Any
8import logging
10from rich.console import Console
11import traceback
13# Rich imports for TUI rendering
14from rich.panel import Panel
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
25from src.ui.ascii_art import display_welcome_screen
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)
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
43# Import the interactive context manager
44from src.interactive_context import InteractiveContext
46# Global reference to TUI instance for service access
47_tui_instance = None
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()
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)
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
70 return Console()
71 return _tui_instance.console
74class ChuckTUI:
75 """
76 Main TUI interface for CHUCK AI.
77 Handles user interaction, execution via ChuckService, and result display.
78 """
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
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
93 def get_service(self):
94 """Get the current ChuckService instance."""
95 return self.service
97 def set_service(self, service):
98 """Set a new ChuckService instance."""
99 self.service = service
100 return True
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"]
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
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]")
120 # Combine all commands and remove duplicates
121 all_commands = builtin_commands + service_commands
122 return sorted(list(set(all_commands)))
124 def _check_first_run(self) -> bool:
125 """
126 Check if this is the first run of the application and configuration is needed.
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
134 config_manager = get_config_manager()
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
148 return False
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")
158 # Display welcome screen
159 display_welcome_screen(self.console)
161 # Initialize the interactive context
162 interactive_context = InteractiveContext()
163 history_file = os.path.expanduser("~/.chuck_history")
165 # Get available commands from service for autocomplete
166 commands = self._get_available_commands()
168 # Set up command completer with slash prefix
169 command_completer = WordCompleter(
170 commands, ignore_case=True, match_middle=True, sentence=True
171 )
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 )
191 # Key bindings to accept history suggestions with Tab
192 bindings = KeyBindings()
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)
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 )
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")
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
232 # Use prompt toolkit with interactive styling
233 prompt_message = HTML(
234 "<interactive-prompt>chuck (interactive) ></interactive-prompt> "
235 )
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
247 user_input = session.prompt(
248 prompt_message,
249 is_password=hide_input,
250 enable_history_search=not hide_input,
251 ).strip()
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
261 # Process the interactive input
262 result = self.service.execute_command(
263 current_cmd, interactive_input=user_input
264 )
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
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
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()
312 # Skip empty commands
313 if not command:
314 continue
316 # Process user command
317 self._process_command(command)
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
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]")
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
355 # Use shlex parsing for commands that clearly need it:
356 # 1. Flag-style arguments (--flag)
357 # 2. Intentional quoted strings (balanced quotes)
359 # Always use shlex for flag-style arguments
360 if "--" in command:
361 return True
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
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
387 return False
389 def _process_command(self, command):
390 """Process a user command in the TUI interface."""
391 interactive_context = InteractiveContext()
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
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]")
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()
420 cmd = parts[0].lower()
421 args = parts[1:]
423 # Process special commands
424 if cmd == "/debug":
425 self._handle_debug(args)
426 return
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)
437 if not result:
438 return
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
448 # Process command result
449 self._process_command_result(cmd, result)
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]")
470 # Skip if no data to display
471 if not result.data:
472 return
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)
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
552 status = "ON" if self.debug else "OFF"
553 self.console.print(
554 f"[{SUCCESS_STYLE}]Debug mode is now {status}[/{SUCCESS_STYLE}]"
555 )
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
562 command_def = get_command(tool_name)
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")
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)
576 except Exception as e:
577 # Handle pagination cancellation specially - let it bubble up
578 from src.exceptions import PaginationCancelled
580 if isinstance(e, PaginationCancelled):
581 raise # Re-raise to bubble up to main TUI loop
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]")
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
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 )
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
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
657 # Extract key information based on tool type
658 status_line = f"[dim cyan]→[/dim cyan] {friendly_name}"
659 metrics = []
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]"
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")
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")
681 # Tag-specific metrics
682 if "tagged_columns" in tool_result:
683 metrics.append(f"{len(tool_result['tagged_columns'])} columns tagged")
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']}")
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']}")
697 # Generic message fallback
698 if not metrics and "message" in tool_result:
699 metrics.append(tool_result["message"])
701 # Format the condensed display
702 if metrics:
703 status_line += f" ({', '.join(metrics)})"
705 self.console.print(status_line)
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}]")
712 if result.error and self.debug:
713 # Use Rich's traceback rendering for better formatting
714 self.console.print_exception(show_locals=True)
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
721 catalogs = data.get("catalogs", [])
722 current_catalog = data.get("current_catalog")
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()
729 # Define column styling based on the active catalog
730 def name_style(name):
731 return "bold green" if name == current_catalog else None
733 style_map = {"name": name_style}
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()
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 )
752 if current_catalog:
753 self.console.print(
754 f"\nCurrent catalog: [{SUCCESS_STYLE}]{current_catalog}[/{SUCCESS_STYLE}]"
755 )
757 # Raise PaginationCancelled to return to chuck > prompt immediately
758 # This prevents agent from continuing processing after catalog display is complete
759 raise PaginationCancelled()
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
766 schemas = data.get("schemas", [])
767 catalog_name = data.get("catalog_name", "")
768 current_schema = data.get("current_schema")
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()
777 # Define column styling based on the active schema
778 def name_style(name):
779 return "bold green" if name == current_schema else None
781 style_map = {"name": name_style}
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 )
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 )
801 # Raise PaginationCancelled to return to chuck > prompt immediately
802 # This prevents agent from continuing processing after schema display is complete
803 raise PaginationCancelled()
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
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))
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()
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
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
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
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)
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
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
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 }
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 )
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 )
907 # Raise PaginationCancelled to return to chuck > prompt immediately
908 # This prevents agent from continuing processing after table display is complete
909 raise PaginationCancelled()
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
916 # Use imported function to get the active model for highlighting
917 active_model = get_active_model()
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()
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 }
935 # Extract and process state information
936 state = model.get("state", {})
937 ready_status = state.get("ready", "UNKNOWN").upper()
938 processed["status"] = ready_status
940 # Add to our list
941 processed_models.append(processed)
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
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
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)"
965 # Set up style map
966 style_map = {"name": name_style, "status": status_style}
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 )
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 )
986 # Raise PaginationCancelled to return to chuck > prompt immediately
987 # This prevents agent from continuing processing after model display is complete
988 raise PaginationCancelled()
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
995 models = data.get("models", [])
996 active_model = data.get("active_model")
997 detailed = data.get("detailed", False)
998 filter_text = data.get("filter")
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()
1010 # Display header with filter information if applicable
1011 title = "Available Models"
1012 if filter_text:
1013 title += f" matching '{filter_text}'"
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 }
1024 # Get state information
1025 state = model.get("state", {})
1026 ready_status = state.get("ready", "UNKNOWN").upper()
1027 processed["status"] = ready_status
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")
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
1042 # Add to our list
1043 processed_models.append(processed)
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
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"
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 }
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"]
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 )
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 )
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()
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
1116 warehouses = data.get("warehouses", [])
1117 current_warehouse_id = data.get("current_warehouse_id")
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()
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)
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
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
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"
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 }
1167 # Set maximum lengths for fields
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 )
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 )
1187 # Raise PaginationCancelled to return to chuck > prompt immediately
1188 # This prevents agent from continuing processing after warehouse display is complete
1189 raise PaginationCancelled()
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
1196 volumes = data.get("volumes", [])
1197 catalog_name = data.get("catalog_name", "")
1198 schema_name = data.get("schema_name", "")
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()
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)
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
1228 # Set up style map
1229 style_map = {"type": type_style}
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 )
1243 # Raise PaginationCancelled to return to chuck > prompt immediately
1244 # This prevents agent from continuing processing after volume display is complete
1245 raise PaginationCancelled()
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
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")
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 ]
1268 # Define styling functions
1269 def value_style(value, row):
1270 setting = row.get("setting", "")
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"
1286 # Set up style map
1287 style_map = {"value": lambda value, row: value_style(value, row)}
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 )
1301 # If permissions data is available, display it
1302 permissions_data = data.get("permissions")
1303 if permissions_data:
1304 self._display_permissions(permissions_data)
1306 def _display_permissions(self, permissions_data: Dict[str, Any]) -> None:
1307 """
1308 Display detailed permission check results.
1310 Args:
1311 permissions_data: Dictionary of permission check results
1312 """
1313 from src.ui.table_formatter import display_table
1315 if not permissions_data:
1316 self.console.print(
1317 f"[{WARNING_STYLE}]No permission data available.[/{WARNING_STYLE}]"
1318 )
1319 return
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")
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)
1343 # Define styling function for status column
1344 def status_style(status, row):
1345 return "green" if row.get("authorized") else "red"
1347 # Set up style map
1348 style_map = {"status": status_style}
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 )
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]")
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
1373 table = data.get("table", {})
1374 full_name = data.get("full_name", "")
1375 has_delta_metadata = data.get("has_delta_metadata", False)
1377 if not table:
1378 self.console.print(
1379 f"[{WARNING_STYLE}]No table details available.[/{WARNING_STYLE}]"
1380 )
1381 return
1383 # Display table header
1384 self.console.print(
1385 f"\n[{TABLE_TITLE_STYLE}]Table Details: {full_name}[/{TABLE_TITLE_STYLE}]"
1386 )
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 ]
1404 for prop, value in properties:
1405 if value: # Only include non-empty values
1406 basic_info.append({"property": prop, "value": value})
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 )
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 )
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 )
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
1453 props_for_display.append({"property": prop, "value": value})
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 )
1465 # Display Delta metadata if available
1466 if has_delta_metadata and "delta" in table:
1467 delta_info = table.get("delta", {})
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 ]
1481 for prop, value in delta_properties:
1482 if value: # Only include non-empty values
1483 delta_for_display.append({"property": prop, "value": value})
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 )
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
1500 catalog = data
1502 if not catalog:
1503 self.console.print(
1504 f"[{WARNING_STYLE}]No catalog details available.[/{WARNING_STYLE}]"
1505 )
1506 return
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 )
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 ]
1529 for prop, value in properties:
1530 if value: # Only include non-empty values
1531 basic_info.append({"property": prop, "value": value})
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 )
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
1546 schema = data
1548 if not schema:
1549 self.console.print(
1550 f"[{WARNING_STYLE}]No schema details available.[/{WARNING_STYLE}]"
1551 )
1552 return
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 )
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 ]
1576 for prop, value in properties:
1577 if value: # Only include non-empty values
1578 basic_info.append({"property": prop, "value": value})
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 )
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
1593 if not data:
1594 self.console.print(
1595 f"[{WARNING_STYLE}]No PII scan results available.[/{WARNING_STYLE}]"
1596 )
1597 return
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)
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 )
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
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)
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 )
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 )
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", [])
1661 if pii_columns:
1662 self.console.print(
1663 f"\n[bold]PII Columns in {table_name}:[/bold]"
1664 )
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 )
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 )
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
1695 if not data:
1696 self.console.print(
1697 f"[{WARNING_STYLE}]No SQL results available.[/{WARNING_STYLE}]"
1698 )
1699 return
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")
1706 if not rows:
1707 self.console.print(
1708 f"[{WARNING_STYLE}]Query returned no results.[/{WARNING_STYLE}]"
1709 )
1710 return
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 )
1718 if should_paginate:
1719 self._display_paginated_sql_results_local(data)
1720 return
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)
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 += ")"
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 )
1749 # Raise PaginationCancelled to return to chuck > prompt immediately
1750 # This prevents agent from continuing processing after SQL display is complete
1751 raise PaginationCancelled()
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)
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
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")
1771 if not external_links:
1772 self.console.print(
1773 f"[{WARNING_STYLE}]No external data links available.[/{WARNING_STYLE}]"
1774 )
1775 return
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 )
1785 rows_displayed = 0
1786 page_num = 1
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
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
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)
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 += ")"
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 )
1836 rows_displayed += len(rows)
1837 page_num += 1
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()
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 )
1852 # Get user input
1853 try:
1854 if sys.stdin.isatty():
1855 import readchar
1857 char = readchar.readchar()
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
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
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 )
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
1916 columns = data.get("columns", [])
1917 rows = data.get("rows", [])
1918 execution_time = data.get("execution_time_ms")
1919 total_rows = len(rows)
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
1927 # Local pagination for large row sets
1928 page_size = 50
1929 current_position = 0
1930 page_num = 1
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]
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)
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 += ")"
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 )
1968 current_position = end_position
1969 page_num += 1
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()
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 )
1984 # Get user input
1985 try:
1986 if sys.stdin.isatty():
1987 import readchar
1989 char = readchar.readchar()
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
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
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 )