diff --git a/databricks-mcp-server/databricks_mcp_server/tools/genie.py b/databricks-mcp-server/databricks_mcp_server/tools/genie.py index 5a606ecf..103a6dae 100644 --- a/databricks-mcp-server/databricks_mcp_server/tools/genie.py +++ b/databricks-mcp-server/databricks_mcp_server/tools/genie.py @@ -42,6 +42,7 @@ def create_or_update_genie( description: Optional[str] = None, sample_questions: Optional[List[str]] = None, space_id: Optional[str] = None, + serialized_space: Optional[str] = None, ) -> Dict[str, Any]: """ Create or update a Genie Space for SQL-based data exploration. @@ -58,6 +59,9 @@ def create_or_update_genie( description: Optional description of what the Genie space does sample_questions: Optional list of sample questions to help users space_id: Optional existing space_id to update instead of create + serialized_space: Optional JSON string containing full space configuration + (settings, instructions). Use this to import/clone a Genie space + exported via get_genie with include_serialized_space=True. Returns: Dictionary with: @@ -75,6 +79,14 @@ def create_or_update_genie( ... sample_questions=["What were total sales last month?"] ... ) {"space_id": "abc123...", "display_name": "Sales Analytics", "operation": "created", ...} + + Clone a space: + >>> source = get_genie(space_id="abc123", include_serialized_space=True) + >>> create_or_update_genie( + ... display_name="Sales Analytics (Copy)", + ... table_identifiers=source["table_identifiers"], + ... serialized_space=source["serialized_space"] + ... ) """ try: description = with_description_footer(description) @@ -99,6 +111,7 @@ def create_or_update_genie( warehouse_id=warehouse_id, table_identifiers=table_identifiers, sample_questions=sample_questions, + serialized_space=serialized_space, ) else: return {"error": f"Genie space {space_id} not found"} @@ -113,6 +126,7 @@ def create_or_update_genie( warehouse_id=warehouse_id, table_identifiers=table_identifiers, sample_questions=sample_questions, + serialized_space=serialized_space, ) space_id = existing.space_id else: @@ -121,6 +135,7 @@ def create_or_update_genie( warehouse_id=warehouse_id, table_identifiers=table_identifiers, description=description, + serialized_space=serialized_space, ) space_id = result.get("space_id", "") @@ -154,7 +169,10 @@ def create_or_update_genie( @mcp.tool -def get_genie(space_id: Optional[str] = None) -> Dict[str, Any]: +def get_genie( + space_id: Optional[str] = None, + include_serialized_space: bool = False, +) -> Dict[str, Any]: """ Get details of a Genie Space, or list all spaces. @@ -163,6 +181,9 @@ def get_genie(space_id: Optional[str] = None) -> Dict[str, Any]: Args: space_id: The Genie space ID. If omitted, lists all spaces. + include_serialized_space: If True, includes the serialized_space field + containing the full space configuration (settings, instructions). + Useful for exporting a space to clone or import elsewhere. Returns: Single space dict (if space_id provided) or {"spaces": [...]}. @@ -173,11 +194,15 @@ def get_genie(space_id: Optional[str] = None) -> Dict[str, Any]: >>> get_genie() {"spaces": [{"space_id": "abc123...", "title": "Sales Analytics", ...}, ...]} + + Export for cloning: + >>> get_genie("abc123...", include_serialized_space=True) + {"space_id": "abc123...", ..., "serialized_space": "{...}"} """ if space_id: try: manager = _get_manager() - result = manager.genie_get(space_id) + result = manager.genie_get(space_id, include_serialized_space=include_serialized_space) if not result: return {"error": f"Genie space {space_id} not found"} @@ -185,7 +210,7 @@ def get_genie(space_id: Optional[str] = None) -> Dict[str, Any]: questions_response = manager.genie_list_questions(space_id, question_type="SAMPLE_QUESTION") sample_questions = [q.get("question_text", "") for q in questions_response.get("curated_questions", [])] - return { + response = { "space_id": result.get("space_id", space_id), "display_name": result.get("display_name", ""), "description": result.get("description", ""), @@ -193,6 +218,11 @@ def get_genie(space_id: Optional[str] = None) -> Dict[str, Any]: "table_identifiers": result.get("table_identifiers", []), "sample_questions": sample_questions, } + + if include_serialized_space and result.get("serialized_space"): + response["serialized_space"] = result["serialized_space"] + + return response except Exception as e: return {"error": f"Failed to get Genie space {space_id}: {e}"} @@ -246,6 +276,92 @@ def delete_genie(space_id: str) -> Dict[str, Any]: return {"success": False, "space_id": space_id, "error": str(e)} +@mcp.tool +def clone_genie( + source_space_id: str, + new_display_name: str, + warehouse_id: Optional[str] = None, + description: Optional[str] = None, +) -> Dict[str, Any]: + """ + Clone a Genie Space by exporting its full configuration and creating a new space. + + Exports the source space (including settings, instructions, sample questions) + and imports it as a new space. Useful for promoting spaces across environments + or creating variants for different teams. + + Args: + source_space_id: The Genie space ID to clone from + new_display_name: Display name for the cloned space + warehouse_id: Optional warehouse ID for the clone. If not provided, + uses the same warehouse as the source space. + description: Optional description for the clone. If not provided, + uses the source space's description. + + Returns: + Dictionary with: + - space_id: The new cloned space ID + - display_name: The new display name + - source_space_id: The original space ID + - operation: 'cloned' + + Example: + >>> clone_genie( + ... source_space_id="abc123...", + ... new_display_name="Sales Analytics (Staging)", + ... ) + {"space_id": "def456...", "display_name": "Sales Analytics (Staging)", ...} + """ + try: + manager = _get_manager() + + source = manager.genie_get(source_space_id, include_serialized_space=True) + if not source: + return {"error": f"Source Genie space {source_space_id} not found"} + + target_warehouse = warehouse_id or source.get("warehouse_id") + if not target_warehouse: + target_warehouse = manager.get_best_warehouse_id() + if not target_warehouse: + return {"error": "No SQL warehouses available. Please provide a warehouse_id."} + + target_description = description or source.get("description", "") + + result = manager.genie_create( + display_name=new_display_name, + warehouse_id=target_warehouse, + table_identifiers=source.get("table_identifiers", []), + description=target_description, + serialized_space=source.get("serialized_space"), + ) + + new_space_id = result.get("space_id", "") + + try: + if new_space_id: + from ..manifest import track_resource + + track_resource( + resource_type="genie_space", + name=new_display_name, + resource_id=new_space_id, + ) + except Exception: + pass + + return { + "space_id": new_space_id, + "display_name": new_display_name, + "source_space_id": source_space_id, + "operation": "cloned", + "warehouse_id": target_warehouse, + "table_count": len(source.get("table_identifiers", [])), + } + + except Exception as e: + return {"error": f"Failed to clone Genie space: {e}"} + + # ============================================================================ # Genie Conversation API Tools # ============================================================================ diff --git a/databricks-tools-core/databricks_tools_core/agent_bricks/manager.py b/databricks-tools-core/databricks_tools_core/agent_bricks/manager.py index 7b829bbc..ac2f4368 100644 --- a/databricks-tools-core/databricks_tools_core/agent_bricks/manager.py +++ b/databricks-tools-core/databricks_tools_core/agent_bricks/manager.py @@ -793,10 +793,21 @@ def mas_list_evaluation_runs( # Genie Space Operations # ======================================================================== - def genie_get(self, space_id: str) -> Optional[GenieSpaceDict]: - """Get Genie space by ID.""" + def genie_get( + self, space_id: str, include_serialized_space: bool = False + ) -> Optional[GenieSpaceDict]: + """Get Genie space by ID. + + Args: + space_id: The Genie space ID + include_serialized_space: If True, includes the serialized_space field + containing full space configuration (settings, instructions, etc.) + """ try: - return self._get(f"/api/2.0/data-rooms/{space_id}") + params = {} + if include_serialized_space: + params["include_serialized_space"] = "true" + return self._get(f"/api/2.0/data-rooms/{space_id}", params=params or None) except Exception as e: if "does not exist" in str(e).lower() or "not found" in str(e).lower(): return None @@ -812,6 +823,7 @@ def genie_create( parent_folder_id: Optional[str] = None, create_dir: bool = True, run_as_type: str = "VIEWER", + serialized_space: Optional[str] = None, ) -> Dict[str, Any]: """Create a Genie space. @@ -824,6 +836,8 @@ def genie_create( parent_folder_id: Optional parent folder ID create_dir: Whether to create parent folder if missing run_as_type: Run as type (default: "VIEWER") + serialized_space: Optional JSON string containing full space configuration + (settings, instructions). Used to import/clone a Genie space. Returns: Created Genie space data @@ -838,6 +852,9 @@ def genie_create( "run_as_type": run_as_type, } + if serialized_space: + room_payload["serialized_space"] = serialized_space + if description: room_payload["description"] = description @@ -869,6 +886,7 @@ def genie_update( warehouse_id: Optional[str] = None, table_identifiers: Optional[List[str]] = None, sample_questions: Optional[List[str]] = None, + serialized_space: Optional[str] = None, ) -> Dict[str, Any]: """Update a Genie space. @@ -879,6 +897,8 @@ def genie_update( warehouse_id: Optional new warehouse ID table_identifiers: Optional new table identifiers sample_questions: Optional sample questions (replaces all existing) + serialized_space: Optional JSON string containing full space configuration + (settings, instructions). Replaces the existing configuration. Returns: Updated Genie space data @@ -913,6 +933,9 @@ def genie_update( if current_space.get(field): update_payload[field] = current_space[field] + if serialized_space: + update_payload["serialized_space"] = serialized_space + result = self._patch(f"/api/2.0/data-rooms/{space_id}", update_payload) if sample_questions is not None: