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
« 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.
4This module contains the handler for creating a Stitch report notebook
5based on a unified table.
6"""
8import logging
9from typing import Optional, Any
11from ..clients.databricks import DatabricksAPIClient
12from ..command_registry import CommandDefinition
13from ..metrics_collector import get_metrics_collector
14from .base import CommandResult
17def handle_command(
18 client: Optional[DatabricksAPIClient], **kwargs: Any
19) -> CommandResult:
20 """
21 Create a Stitch report notebook for a unified table.
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")
33 # If rest is provided, use it as the notebook name (overrides name parameter)
34 if rest_args:
35 custom_name = rest_args
37 if not client:
38 return CommandResult(
39 False, message="Client is required for creating Stitch report notebooks."
40 )
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 )
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 )
57 # Get metrics collector
58 metrics_collector = get_metrics_collector()
60 try:
61 # Call the create_stitch_notebook function with optional name
62 result = client.create_stitch_notebook(table_path, custom_name)
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 )
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
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 )
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)