Coverage for src/chuck_data/commands/setup_stitch.py: 0%

156 statements  

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

1""" 

2Command handler for Stitch integration setup. 

3 

4This module contains the handler for setting up a Stitch integration by scanning 

5for PII columns and creating a configuration file. 

6""" 

7 

8import logging 

9from typing import Optional 

10 

11from ..clients.databricks import DatabricksAPIClient 

12from ..llm.client import LLMClient 

13from ..command_registry import CommandDefinition 

14from ..config import get_active_catalog, get_active_schema 

15from ..metrics_collector import get_metrics_collector 

16from ..interactive_context import InteractiveContext 

17from ..ui.theme import SUCCESS_STYLE, ERROR_STYLE, INFO_STYLE, WARNING 

18from ..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) 

26 

27 

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}]") 

31 

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']}") 

36 

37 # Show tables and fields 

38 table_count = len(stitch_config["tables"]) 

39 console.print(f"• Tables to process: {table_count}") 

40 

41 total_fields = sum(len(table["fields"]) for table in stitch_config["tables"]) 

42 console.print(f"• Total PII fields: {total_fields}") 

43 

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)") 

49 

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']}") 

57 

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 ) 

64 

65 

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") 

74 

75 

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. 

84 

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") 

95 

96 if not client: 

97 return CommandResult(False, message="Client is required for Stitch setup.") 

98 

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) 

103 

104 # Interactive mode - use context management 

105 context = InteractiveContext() 

106 console = get_console() 

107 

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 ) 

114 

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 ) 

122 

123 current_phase = builder_data.get("phase", "review") 

124 

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 ) 

134 

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 ) 

142 

143 

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() 

153 

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 ) 

159 

160 # Create a LLM client instance to pass to the helper 

161 llm_client = LLMClient() 

162 

163 # Get metrics collector 

164 metrics_collector = get_metrics_collector() 

165 

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 ) 

189 

190 return CommandResult(False, message=prep_result["error"], data=prep_result) 

191 

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 ) 

215 

216 return CommandResult( 

217 False, message=stitch_result_data["error"], data=stitch_result_data 

218 ) 

219 

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 ) 

235 

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 ) 

246 

247 

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() 

258 

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 ) 

264 

265 # Set context as active for interactive mode 

266 context.set_active_context("setup-stitch") 

267 

268 console.print( 

269 f"\n[{INFO_STYLE}]Preparing Stitch configuration for {target_catalog}.{target_schema}...[/{INFO_STYLE}]" 

270 ) 

271 

272 # Create LLM client 

273 llm_client = LLMClient() 

274 

275 # Prepare the configuration 

276 prep_result = _helper_prepare_stitch_config( 

277 client, llm_client, target_catalog, target_schema 

278 ) 

279 

280 if prep_result.get("error"): 

281 context.clear_active_context("setup-stitch") 

282 return CommandResult(False, message=prep_result["error"]) 

283 

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 

291 

292 # Display the configuration preview 

293 _display_config_preview( 

294 console, prep_result["stitch_config"], prep_result["metadata"] 

295 ) 

296 _display_confirmation_prompt(console) 

297 

298 return CommandResult( 

299 True, message="" # Empty message - let the console output speak for itself 

300 ) 

301 

302 

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 

311 

312 user_input_lower = user_input.lower().strip() 

313 

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 ) 

324 

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.") 

330 

331 # Otherwise, treat as modification request 

332 console.print( 

333 f"\n[{INFO_STYLE}]Modifying configuration based on your request...[/{INFO_STYLE}]" 

334 ) 

335 

336 modify_result = _helper_modify_stitch_config( 

337 stitch_config, user_input, llm_client, metadata 

338 ) 

339 

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 ) 

351 

352 # Update stored config 

353 updated_config = modify_result["stitch_config"] 

354 context.store_context_data("setup-stitch", "stitch_config", updated_config) 

355 

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"]) 

359 

360 # Show updated preview 

361 _display_config_preview(console, updated_config, metadata) 

362 _display_confirmation_prompt(console) 

363 

364 return CommandResult( 

365 True, 

366 message="Please review the updated configuration and choose: 'launch', more changes, or 'cancel'.", 

367 ) 

368 

369 

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"] 

377 

378 user_input_lower = user_input.lower().strip() 

379 

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}]") 

390 

391 # Launch the job 

392 launch_result = _helper_launch_stitch_job(client, stitch_config, metadata) 

393 

394 # Clear context after launch (success or failure) 

395 context.clear_active_context("setup-stitch") 

396 

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 ) 

420 

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 ) 

440 

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 ) 

449 

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.") 

454 

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 ) 

462 

463 

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)