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

32 statements  

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

1""" 

2Command handler for adding a Stitch Report notebook. 

3 

4This module contains the handler for creating a Stitch report notebook 

5based on a unified table. 

6""" 

7 

8import logging 

9from typing import Optional, Any 

10 

11from ..clients.databricks import DatabricksAPIClient 

12from ..command_registry import CommandDefinition 

13from ..metrics_collector import get_metrics_collector 

14from .base import CommandResult 

15 

16 

17def handle_command( 

18 client: Optional[DatabricksAPIClient], **kwargs: Any 

19) -> CommandResult: 

20 """ 

21 Create a Stitch report notebook for a unified table. 

22 

23 Args: 

24 client: API client instance 

25 **kwargs: 

26 table_path (str): Fully qualified table path in format catalog.schema.table 

27 name (str, optional): Optional custom name for the notebook 

28 """ 

29 table_path: Optional[str] = kwargs.get("table_path") 

30 custom_name: Optional[str] = kwargs.get("name") 

31 rest_args: Optional[str] = kwargs.get("rest") 

32 

33 # If rest is provided, use it as the notebook name (overrides name parameter) 

34 if rest_args: 

35 custom_name = rest_args 

36 

37 if not client: 

38 return CommandResult( 

39 False, message="Client is required for creating Stitch report notebooks." 

40 ) 

41 

42 try: 

43 if not table_path: 

44 return CommandResult( 

45 False, 

46 message="Table path must be provided in format 'catalog.schema.table'.", 

47 ) 

48 

49 # Validate the table path format 

50 path_parts = table_path.split(".") 

51 if len(path_parts) != 3: 

52 return CommandResult( 

53 False, 

54 message="Table path must be in the format 'catalog.schema.table'.", 

55 ) 

56 

57 # Get metrics collector 

58 metrics_collector = get_metrics_collector() 

59 

60 try: 

61 # Call the create_stitch_notebook function with optional name 

62 result = client.create_stitch_notebook(table_path, custom_name) 

63 

64 # Track successful event 

65 metrics_collector.track_event( 

66 prompt="add-stitch-report command", 

67 tools=[ 

68 { 

69 "name": "add-stitch-report", 

70 "arguments": { 

71 "table_path": table_path, 

72 "notebook_name": custom_name, 

73 }, 

74 } 

75 ], 

76 additional_data={ 

77 "event_context": "direct_add-stitch-report_command", 

78 "status": "success", 

79 "notebook_path": result.get("path", ""), 

80 }, 

81 ) 

82 

83 return CommandResult( 

84 True, 

85 data=result, 

86 message=f"Successfully created Stitch report notebook at {result.get('path')}", 

87 ) 

88 except Exception as e: 

89 # Track error event 

90 metrics_collector.track_event( 

91 prompt="add-stitch-report command", 

92 tools=[ 

93 { 

94 "name": "add-stitch-report", 

95 "arguments": { 

96 "table_path": table_path, 

97 "notebook_name": custom_name, 

98 }, 

99 } 

100 ], 

101 error=str(e), 

102 additional_data={ 

103 "event_context": "direct_add-stitch-report_command", 

104 "status": "error", 

105 }, 

106 ) 

107 raise 

108 

109 except Exception as e: 

110 logging.error(f"Error creating Stitch report: {e}", exc_info=True) 

111 return CommandResult( 

112 False, error=e, message=f"Error creating Stitch report: {str(e)}" 

113 ) 

114 

115 

116DEFINITION = CommandDefinition( 

117 name="add-stitch-report", 

118 description="Create a Stitch report notebook for a unified table", 

119 handler=handle_command, 

120 parameters={ 

121 "table_path": { 

122 "type": "string", 

123 "description": "Fully qualified table path in format 'catalog.schema.table'", 

124 }, 

125 "name": { 

126 "type": "string", 

127 "description": "Optional: Custom name for the notebook", 

128 }, 

129 "rest": { 

130 "type": "string", 

131 "description": "Additional arguments will be treated as part of the notebook name", 

132 }, 

133 }, 

134 required_params=["table_path"], 

135 tui_aliases=["/add-stitch-report"], 

136 visible_to_user=True, 

137 visible_to_agent=True, 

138 needs_api_client=True, 

139 usage_hint="/add-stitch-report catalog.schema.table [notebook name]", 

140)