Coverage for src/commands/setup_stitch.py: 37%
156 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
« prev ^ index » next coverage.py v7.8.0, created at 2025-06-05 22:56 -0700
1"""
2Command handler for Stitch integration setup.
4This module contains the handler for setting up a Stitch integration by scanning
5for PII columns and creating a configuration file.
6"""
8import logging
9from typing import Optional
11from src.clients.databricks import DatabricksAPIClient
12from src.llm.client import LLMClient
13from src.command_registry import CommandDefinition
14from src.config import get_active_catalog, get_active_schema
15from src.metrics_collector import get_metrics_collector
16from src.interactive_context import InteractiveContext
17from src.ui.theme import SUCCESS_STYLE, ERROR_STYLE, INFO_STYLE, WARNING
18from src.ui.tui import get_console
19from .base import CommandResult
20from .stitch_tools import (
21 _helper_setup_stitch_logic,
22 _helper_prepare_stitch_config,
23 _helper_modify_stitch_config,
24 _helper_launch_stitch_job,
25)
28def _display_config_preview(console, stitch_config, metadata):
29 """Display a preview of the Stitch configuration to the user."""
30 console.print(f"\n[{INFO_STYLE}]Stitch Configuration Preview:[/{INFO_STYLE}]")
32 # Show basic info
33 console.print(f"• Target: {metadata['target_catalog']}.{metadata['target_schema']}")
34 console.print(f"• Job Name: {metadata['stitch_job_name']}")
35 console.print(f"• Config Path: {metadata['config_file_path']}")
37 # Show tables and fields
38 table_count = len(stitch_config["tables"])
39 console.print(f"• Tables to process: {table_count}")
41 total_fields = sum(len(table["fields"]) for table in stitch_config["tables"])
42 console.print(f"• Total PII fields: {total_fields}")
44 if table_count > 0:
45 console.print("\nTables:")
46 for table in stitch_config["tables"]:
47 field_count = len(table["fields"])
48 console.print(f" - {table['path']} ({field_count} fields)")
50 # Show all fields
51 for field in table["fields"]:
52 semantics = ", ".join(field.get("semantics", []))
53 if semantics:
54 console.print(f" • {field['field-name']} ({semantics})")
55 else:
56 console.print(f" • {field['field-name']}")
58 # Show unsupported columns if any
59 unsupported = metadata.get("unsupported_columns", [])
60 if unsupported:
61 console.print(
62 f"\n[{WARNING}]Note: {sum(len(t['columns']) for t in unsupported)} columns excluded due to unsupported types[/{WARNING}]"
63 )
66def _display_confirmation_prompt(console):
67 """Display the confirmation prompt to the user."""
68 console.print(f"\n[{INFO_STYLE}]What would you like to do?[/{INFO_STYLE}]")
69 console.print("• Type 'launch' or 'yes' to launch the job")
70 console.print(
71 "• Describe changes (e.g., 'remove table X', 'add email semantic to field Y')"
72 )
73 console.print("• Type 'cancel' to abort the setup")
76def handle_command(
77 client: Optional[DatabricksAPIClient],
78 interactive_input: str = None,
79 auto_confirm: bool = False,
80 **kwargs,
81) -> CommandResult:
82 """
83 Set up a Stitch integration with interactive configuration review.
85 Args:
86 client: API client instance
87 interactive_input: User input for interactive mode
88 auto_confirm: Skip confirmation and launch immediately
89 **kwargs:
90 catalog_name (str, optional): Target catalog name
91 schema_name (str, optional): Target schema name
92 """
93 catalog_name_arg: Optional[str] = kwargs.get("catalog_name")
94 schema_name_arg: Optional[str] = kwargs.get("schema_name")
96 if not client:
97 return CommandResult(False, message="Client is required for Stitch setup.")
99 # Determine if legacy auto-confirm mode was explicitly requested
100 explicit_auto_confirm = auto_confirm or kwargs.get("auto_confirm") is True
101 if explicit_auto_confirm:
102 return _handle_legacy_setup(client, catalog_name_arg, schema_name_arg)
104 # Interactive mode - use context management
105 context = InteractiveContext()
106 console = get_console()
108 try:
109 # Phase determination
110 if not interactive_input: # First call - Phase 1: Prepare config
111 return _phase_1_prepare_config(
112 client, context, console, catalog_name_arg, schema_name_arg
113 )
115 # Get stored context data
116 builder_data = context.get_context_data("setup-stitch")
117 if not builder_data:
118 return CommandResult(
119 False,
120 message="Stitch setup context lost. Please run /setup-stitch again.",
121 )
123 current_phase = builder_data.get("phase", "review")
125 if current_phase == "review":
126 return _phase_2_handle_review(client, context, console, interactive_input)
127 elif current_phase == "ready_to_launch":
128 return _phase_3_launch_job(client, context, console, interactive_input)
129 else:
130 return CommandResult(
131 False,
132 message=f"Unknown phase: {current_phase}. Please run /setup-stitch again.",
133 )
135 except Exception as e:
136 # Clear context on error
137 context.clear_active_context("setup-stitch")
138 logging.error(f"Stitch setup error: {e}", exc_info=True)
139 return CommandResult(
140 False, error=e, message=f"Error setting up Stitch: {str(e)}"
141 )
144def _handle_legacy_setup(
145 client: DatabricksAPIClient,
146 catalog_name_arg: Optional[str],
147 schema_name_arg: Optional[str],
148) -> CommandResult:
149 """Handle auto-confirm mode using the legacy direct setup approach."""
150 try:
151 target_catalog = catalog_name_arg or get_active_catalog()
152 target_schema = schema_name_arg or get_active_schema()
154 if not target_catalog or not target_schema:
155 return CommandResult(
156 False,
157 message="Target catalog and schema must be specified or active for Stitch setup.",
158 )
160 # Create a LLM client instance to pass to the helper
161 llm_client = LLMClient()
163 # Get metrics collector
164 metrics_collector = get_metrics_collector()
166 # Get the prepared configuration (doesn't launch job anymore)
167 prep_result = _helper_setup_stitch_logic(
168 client, llm_client, target_catalog, target_schema
169 )
170 if prep_result.get("error"):
171 # Track error event
172 metrics_collector.track_event(
173 prompt="setup-stitch command",
174 tools=[
175 {
176 "name": "setup_stitch",
177 "arguments": {
178 "catalog": target_catalog,
179 "schema": target_schema,
180 },
181 }
182 ],
183 error=prep_result.get("error"),
184 additional_data={
185 "event_context": "direct_stitch_command",
186 "status": "error",
187 },
188 )
190 return CommandResult(False, message=prep_result["error"], data=prep_result)
192 # Now we need to explicitly launch the job since _helper_setup_stitch_logic no longer does it
193 stitch_result_data = _helper_launch_stitch_job(
194 client, prep_result["stitch_config"], prep_result["metadata"]
195 )
196 if stitch_result_data.get("error"):
197 # Track error event for launch failure
198 metrics_collector.track_event(
199 prompt="setup_stitch command",
200 tools=[
201 {
202 "name": "setup_stitch",
203 "arguments": {
204 "catalog": target_catalog,
205 "schema": target_schema,
206 },
207 }
208 ],
209 error=stitch_result_data.get("error"),
210 additional_data={
211 "event_context": "direct_stitch_command",
212 "status": "launch_error",
213 },
214 )
216 return CommandResult(
217 False, message=stitch_result_data["error"], data=stitch_result_data
218 )
220 # Track successful stitch setup event
221 metrics_collector.track_event(
222 prompt="setup-stitch command",
223 tools=[
224 {
225 "name": "setup_stitch",
226 "arguments": {"catalog": target_catalog, "schema": target_schema},
227 }
228 ],
229 additional_data={
230 "event_context": "direct_stitch_command",
231 "status": "success",
232 **{k: v for k, v in stitch_result_data.items() if k != "message"},
233 },
234 )
236 return CommandResult(
237 True,
238 data=stitch_result_data,
239 message=stitch_result_data.get("message", "Stitch setup completed."),
240 )
241 except Exception as e:
242 logging.error(f"Legacy stitch setup error: {e}", exc_info=True)
243 return CommandResult(
244 False, error=e, message=f"Error setting up Stitch: {str(e)}"
245 )
248def _phase_1_prepare_config(
249 client: DatabricksAPIClient,
250 context: InteractiveContext,
251 console,
252 catalog_name_arg: Optional[str],
253 schema_name_arg: Optional[str],
254) -> CommandResult:
255 """Phase 1: Prepare the Stitch configuration."""
256 target_catalog = catalog_name_arg or get_active_catalog()
257 target_schema = schema_name_arg or get_active_schema()
259 if not target_catalog or not target_schema:
260 return CommandResult(
261 False,
262 message="Target catalog and schema must be specified or active for Stitch setup.",
263 )
265 # Set context as active for interactive mode
266 context.set_active_context("setup-stitch")
268 console.print(
269 f"\n[{INFO_STYLE}]Preparing Stitch configuration for {target_catalog}.{target_schema}...[/{INFO_STYLE}]"
270 )
272 # Create LLM client
273 llm_client = LLMClient()
275 # Prepare the configuration
276 prep_result = _helper_prepare_stitch_config(
277 client, llm_client, target_catalog, target_schema
278 )
280 if prep_result.get("error"):
281 context.clear_active_context("setup-stitch")
282 return CommandResult(False, message=prep_result["error"])
284 # Store the prepared data in context (don't store llm_client object)
285 context.store_context_data("setup-stitch", "phase", "review")
286 context.store_context_data(
287 "setup-stitch", "stitch_config", prep_result["stitch_config"]
288 )
289 context.store_context_data("setup-stitch", "metadata", prep_result["metadata"])
290 # Note: We'll recreate LLMClient in each phase instead of storing it
292 # Display the configuration preview
293 _display_config_preview(
294 console, prep_result["stitch_config"], prep_result["metadata"]
295 )
296 _display_confirmation_prompt(console)
298 return CommandResult(
299 True, message="" # Empty message - let the console output speak for itself
300 )
303def _phase_2_handle_review(
304 client: DatabricksAPIClient, context: InteractiveContext, console, user_input: str
305) -> CommandResult:
306 """Phase 2: Handle user review and potential config modifications."""
307 builder_data = context.get_context_data("setup-stitch")
308 stitch_config = builder_data["stitch_config"]
309 metadata = builder_data["metadata"]
310 llm_client = LLMClient() # Recreate instead of getting from context
312 user_input_lower = user_input.lower().strip()
314 # Check for launch commands
315 if user_input_lower in ["launch", "yes", "y", "launch it", "go", "proceed"]:
316 # Move to launch phase
317 context.store_context_data("setup-stitch", "phase", "ready_to_launch")
318 console.print(
319 f"\n[{WARNING}]Ready to launch Stitch job. Type 'confirm' to proceed or 'cancel' to abort.[/{WARNING}]"
320 )
321 return CommandResult(
322 True, message="Ready to launch. Type 'confirm' to proceed with job launch."
323 )
325 # Check for cancel
326 if user_input_lower in ["cancel", "abort", "stop", "exit", "quit", "no"]:
327 context.clear_active_context("setup-stitch")
328 console.print(f"\n[{INFO_STYLE}]Stitch setup cancelled.[/{INFO_STYLE}]")
329 return CommandResult(True, message="Stitch setup cancelled.")
331 # Otherwise, treat as modification request
332 console.print(
333 f"\n[{INFO_STYLE}]Modifying configuration based on your request...[/{INFO_STYLE}]"
334 )
336 modify_result = _helper_modify_stitch_config(
337 stitch_config, user_input, llm_client, metadata
338 )
340 if modify_result.get("error"):
341 console.print(
342 f"\n[{ERROR_STYLE}]Error modifying configuration: {modify_result['error']}[/{ERROR_STYLE}]"
343 )
344 console.print(
345 "Please try rephrasing your request or type 'launch' to proceed with current config."
346 )
347 return CommandResult(
348 True,
349 message="Please try rephrasing your request or type 'launch' to proceed.",
350 )
352 # Update stored config
353 updated_config = modify_result["stitch_config"]
354 context.store_context_data("setup-stitch", "stitch_config", updated_config)
356 console.print(f"\n[{SUCCESS_STYLE}]Configuration updated![/{SUCCESS_STYLE}]")
357 if modify_result.get("modification_summary"):
358 console.print(modify_result["modification_summary"])
360 # Show updated preview
361 _display_config_preview(console, updated_config, metadata)
362 _display_confirmation_prompt(console)
364 return CommandResult(
365 True,
366 message="Please review the updated configuration and choose: 'launch', more changes, or 'cancel'.",
367 )
370def _phase_3_launch_job(
371 client: DatabricksAPIClient, context: InteractiveContext, console, user_input: str
372) -> CommandResult:
373 """Phase 3: Final confirmation and job launch."""
374 builder_data = context.get_context_data("setup-stitch")
375 stitch_config = builder_data["stitch_config"]
376 metadata = builder_data["metadata"]
378 user_input_lower = user_input.lower().strip()
380 if user_input_lower in [
381 "confirm",
382 "yes",
383 "y",
384 "launch",
385 "proceed",
386 "go",
387 "make it so",
388 ]:
389 console.print(f"\n[{INFO_STYLE}]Launching Stitch job...[/{INFO_STYLE}]")
391 # Launch the job
392 launch_result = _helper_launch_stitch_job(client, stitch_config, metadata)
394 # Clear context after launch (success or failure)
395 context.clear_active_context("setup-stitch")
397 if launch_result.get("error"):
398 # Track error event
399 metrics_collector = get_metrics_collector()
400 metrics_collector.track_event(
401 prompt="setup-stitch command",
402 tools=[
403 {
404 "name": "setup_stitch",
405 "arguments": {
406 "catalog": metadata["target_catalog"],
407 "schema": metadata["target_schema"],
408 },
409 }
410 ],
411 error=launch_result.get("error"),
412 additional_data={
413 "event_context": "interactive_stitch_command",
414 "status": "error",
415 },
416 )
417 return CommandResult(
418 False, message=launch_result["error"], data=launch_result
419 )
421 # Track successful launch
422 metrics_collector = get_metrics_collector()
423 metrics_collector.track_event(
424 prompt="setup-stitch command",
425 tools=[
426 {
427 "name": "setup_stitch",
428 "arguments": {
429 "catalog": metadata["target_catalog"],
430 "schema": metadata["target_schema"],
431 },
432 }
433 ],
434 additional_data={
435 "event_context": "interactive_stitch_command",
436 "status": "success",
437 **{k: v for k, v in launch_result.items() if k != "message"},
438 },
439 )
441 console.print(
442 f"\n[{SUCCESS_STYLE}]Stitch job launched successfully![/{SUCCESS_STYLE}]"
443 )
444 return CommandResult(
445 True,
446 data=launch_result,
447 message=launch_result.get("message", "Stitch setup completed."),
448 )
450 elif user_input_lower in ["cancel", "abort", "stop", "no"]:
451 context.clear_active_context("setup-stitch")
452 console.print(f"\n[{INFO_STYLE}]Stitch job launch cancelled.[/{INFO_STYLE}]")
453 return CommandResult(True, message="Stitch job launch cancelled.")
455 else:
456 console.print(
457 f"\n[{WARNING}]Please type 'confirm' to launch the job or 'cancel' to abort.[/{WARNING}]"
458 )
459 return CommandResult(
460 True, message="Please type 'confirm' to launch or 'cancel' to abort."
461 )
464DEFINITION = CommandDefinition(
465 name="setup-stitch",
466 description="Interactively set up a Stitch integration with configuration review and modification",
467 handler=handle_command,
468 parameters={
469 "catalog_name": {
470 "type": "string",
471 "description": "Optional: Name of the catalog. If not provided, uses the active catalog",
472 },
473 "schema_name": {
474 "type": "string",
475 "description": "Optional: Name of the schema. If not provided, uses the active schema",
476 },
477 "auto_confirm": {
478 "type": "boolean",
479 "description": "Optional: Skip interactive confirmation and launch job immediately (default: false)",
480 },
481 },
482 required_params=[],
483 tui_aliases=["/setup-stitch"],
484 visible_to_user=True,
485 visible_to_agent=True,
486 supports_interactive_input=True,
487 usage_hint="Example: /setup-stitch or /setup-stitch --auto-confirm to skip confirmation",
488 condensed_action="Setting up Stitch integration",
489)