Refactor create_or_update_genie function for improved handling of serialized_space

- Streamline logic for updating or creating Genie spaces based on the presence of serialized_space and space_id.
- Enhance error handling for non-existent spaces and ensure proper updates are made when serialized_space is provided.
- Update documentation to clarify the workflow for creating and updating Genie spaces, including handling sample questions.

Author: BlitzBricksterYY-db

databricks-mcp-server/databricks_mcp_server/tools/genie.py

@@ -107,104 +107,106 @@ def create_or_update_genie(
 
         operation = "created"
 
-    # When serialized_space is provided, use the public genie/spaces API
-    if serialized_space:
-        if space_id:
-            # Update existing space with serialized config
-            manager.genie_update_with_serialized_space(
-                space_id=space_id,
-                serialized_space=serialized_space,
-                title=display_name,
-                description=description,
-                warehouse_id=warehouse_id,
-            )
-            operation = "updated"
-        else:
-            # Check if exists by name, then create or update
-            existing = manager.genie_find_by_name(display_name)
-            if existing:
-                operation = "updated"
-                space_id = existing.space_id
+        # When serialized_space is provided
+        if serialized_space:
+            if space_id:
+                # Update existing space with serialized config
                 manager.genie_update_with_serialized_space(
                     space_id=space_id,
                     serialized_space=serialized_space,
                     title=display_name,
                     description=description,
                     warehouse_id=warehouse_id,
                 )
-            else:
-                result = manager.genie_import(
-                    warehouse_id=warehouse_id,
-                    serialized_space=serialized_space,
-                    title=display_name,
-                    description=description,
-                )
-                space_id = result.get("space_id", "")
-    else:
-        if space_id:
-            # Update existing space by ID
-            existing = manager.genie_get(space_id)
-            if existing:
                 operation = "updated"
-                manager.genie_update(
-                    space_id=space_id,
-                    display_name=display_name,
-                    description=description,
-                    warehouse_id=warehouse_id,
-                    table_identifiers=table_identifiers,
-                    sample_questions=sample_questions,
-                )
             else:
-                return {"error": f"Genie space {space_id} not found"}
+                # Check if exists by name, then create or update
+                existing = manager.genie_find_by_name(display_name)
+                if existing:
+                    operation = "updated"
+                    space_id = existing.space_id
+                    manager.genie_update_with_serialized_space(
+                        space_id=space_id,
+                        serialized_space=serialized_space,
+                        title=display_name,
+                        description=description,
+                        warehouse_id=warehouse_id,
+                    )
+                else:
+                    result = manager.genie_import(
+                        warehouse_id=warehouse_id,
+                        serialized_space=serialized_space,
+                        title=display_name,
+                        description=description,
+                    )
+                    space_id = result.get("space_id", "")
+        
+        # When serialized_space is not provided
         else:
-            # Check if exists by name first
-            existing = manager.genie_find_by_name(display_name)
-            if existing:
-                operation = "updated"
-                manager.genie_update(
-                    space_id=existing.space_id,
-                    display_name=display_name,
-                    description=description,
-                    warehouse_id=warehouse_id,
-                    table_identifiers=table_identifiers,
-                    sample_questions=sample_questions,
-                )
-                space_id = existing.space_id
+            if space_id:
+                # Update existing space by ID
+                existing = manager.genie_get(space_id)
+                if existing:
+                    operation = "updated"
+                    manager.genie_update(
+                        space_id=space_id,
+                        display_name=display_name,
+                        description=description,
+                        warehouse_id=warehouse_id,
+                        table_identifiers=table_identifiers,
+                        sample_questions=sample_questions,
+                    )
+                else:
+                    return {"error": f"Genie space {space_id} not found"}
             else:
-                # Create new
-                result = manager.genie_create(
-                    display_name=display_name,
-                    warehouse_id=warehouse_id,
-                    table_identifiers=table_identifiers,
-                    description=description,
-                )
-                space_id = result.get("space_id", "")
+                # Check if exists by name first
+                existing = manager.genie_find_by_name(display_name)
+                if existing:
+                    operation = "updated"
+                    manager.genie_update(
+                        space_id=existing.space_id,
+                        display_name=display_name,
+                        description=description,
+                        warehouse_id=warehouse_id,
+                        table_identifiers=table_identifiers,
+                        sample_questions=sample_questions,
+                    )
+                    space_id = existing.space_id
+                else:
+                    # Create new
+                    result = manager.genie_create(
+                        display_name=display_name,
+                        warehouse_id=warehouse_id,
+                        table_identifiers=table_identifiers,
+                        description=description,
+                    )
+                    space_id = result.get("space_id", "")
+
+                    # Add sample questions if provided
+                    if sample_questions and space_id:
+                        manager.genie_add_sample_questions_batch(space_id, sample_questions)
 
-                # Add sample questions if provided
-                if sample_questions and space_id:
-                    manager.genie_add_sample_questions_batch(space_id, sample_questions)
-
-        response = {
-            "space_id": space_id,
-            "display_name": display_name,
-            "operation": operation,
-            "warehouse_id": warehouse_id,
-            "table_count": len(table_identifiers),
-        }
-
-        try:
-            if space_id:
-                from ..manifest import track_resource
+            response = {
+                "space_id": space_id,
+                "display_name": display_name,
+                "operation": operation,
+                "warehouse_id": warehouse_id,
+                "table_count": len(table_identifiers),
+            }
 
-                track_resource(
-                    resource_type="genie_space",
-                    name=display_name,
-                    resource_id=space_id,
-                )
-        except Exception:
-            pass
+            try:
+                if space_id:
+                    from ..manifest import track_resource
+
+                    track_resource(
+                        resource_type="genie_space",
+                        name=display_name,
+                        resource_id=space_id,
+                    )
+            except Exception:
+                pass
 
-        return response
+            return response
 
     except Exception as e:
         return {"error": f"Failed to create/update Genie space '{display_name}': {e}"}

databricks-skills/databricks-genie/SKILL.md

@@ -113,28 +113,92 @@ import_genie(
 
 #### Example: Migrating Genie Spaces from Prod to Dev
 
-When migrating Genie Spaces between environments (e.g., from a `prod` target to a `dev` target defined in your `databricks.yml`), you must update the catalog references within the serialized space. 
+When migrating Genie Spaces between environments (e.g., from a `prod` target to a `dev` target defined in your `databricks.yml`), you must update the catalog references within the serialized space.
 
 **Note:** Genie Space migration assumes that the underlying data assets (schemas and tables) remain structurally identical across environments. The migration of the actual catalogs, schemas, or tables themselves is outside the scope of Genie Space migration skills.
 
-For instance, if your production tables reside in the `healthverity_claims_sample_patient_dataset` catalog, but your development tables are in `healthverity_claims_sample_patient_dataset_dev`, you can perform a string replacement on the exported configuration before importing it into the target workspace:
+##### The Challenge: MCP Servers Are Workspace-Scoped
+
+Each Databricks MCP server instance connects to exactly one workspace (set via `DATABRICKS_CONFIG_PROFILE` at startup). This means a single MCP server cannot export from PROD and import into DEV in the same session — you need two server instances.
+
+##### Recommended Setup: Dual MCP Server Profiles
+
+Configure two Databricks MCP server entries in your IDE's MCP config (e.g. `~/.cursor/mcp.json`), one per workspace:
+
+```json
+"databricks-prod": {
+  "command": "/path/to/.venv/bin/python",
+  "args": ["/path/to/databricks-mcp-server/run_server.py"],
+  "env": { "DATABRICKS_CONFIG_PROFILE": "prod" }
+},
+"databricks-dev": {
+  "command": "/path/to/.venv/bin/python",
+  "args": ["/path/to/databricks-mcp-server/run_server.py"],
+  "env": { "DATABRICKS_CONFIG_PROFILE": "dev" }
+}
+```
+
+Both servers run simultaneously after one IDE reload. This lets you call `export_genie` against `databricks-prod` and `import_genie` against `databricks-dev` within the same conversation — no further reloads needed.
+
+> **Tip:** The Databricks CLI profiles (`prod`, `dev`) referenced above must be defined in `~/.databrickscfg`. Both token-based and OAuth (`auth_type = databricks-cli`) profiles are supported.
+
+##### Full Migration Workflow
+
+**Step 1 — Export from PROD** using the `databricks-prod` MCP server:
 
 ```python
-# 1. Export the Genie Space from the production workspace
+# Call export_genie via the prod-scoped MCP server
 exported = export_genie(space_id="<prod_space_id>")
+# exported["serialized_space"] contains the full config
+# exported["warehouse_id"] is the PROD warehouse — do NOT reuse it for DEV
+```
+
+**Step 2 — Find the DEV warehouse ID:**
 
-# 2. Remap the catalog name for the development environment
+```python
+# Call list_warehouses via the dev-scoped MCP server
+list_warehouses()  # note the warehouse_id for the DEV workspace
+```
+
+**Step 3 — Remap the catalog and import into DEV** using the `databricks-dev` MCP server:
+
+```python
+# Catalog name differs between environments — replace ALL occurrences.
+# serialized_space embeds the catalog in table identifiers, SQL FROM clauses,
+# join specs, and filter snippets, so a single string replace covers everything.
 dev_serialized_space = exported["serialized_space"].replace(
-    "healthverity_claims_sample_patient_dataset", 
-    "healthverity_claims_sample_patient_dataset_dev"
+    "my_prod_catalog",
+    "my_dev_catalog"
 )
 
-# 3. Import the modified space into the dev workspace
-import_genie(
+# Call import_genie via the dev-scoped MCP server
+result = import_genie(
     warehouse_id="<dev_warehouse_id>",
     serialized_space=dev_serialized_space,
-    title="HealthVerity Claims (Dev)"
+    title="My Space"
 )
+# result["space_id"] is the new DEV space ID
+```
+
+**Step 4 — Update `databricks.yml`** with the new DEV space IDs so they are tracked in the bundle:
+
+```yaml
+targets:
+  dev:
+    variables:
+      genie_space_ids: "<new_dev_space_id_1>,<new_dev_space_id_2>,<new_dev_space_id_3>"
+```
+
+**Step 5 — Save exports locally** for version control and future re-migrations:
+
+```json
+// genie_exports/MySpace.json
+{
+  "space_id": "<prod_space_id>",
+  "title": "MySpace",
+  "warehouse_id": "<prod_warehouse_id>",
+  "serialized_space": "{ ... }"
+}
 ```
 
 ## Workflow
@@ -176,7 +240,9 @@ Use these skills in sequence:
 | **`import_genie` fails with permission error** | Ensure you have CREATE privileges in the target workspace folder |
 | **Tables not found after migration** | Catalog name was not remapped — replace the source catalog name in `serialized_space` before calling `import_genie` |
 | **Catalog name appears in SQL queries too** | `serialized_space` embeds catalog in table identifiers, SQL FROM clauses, join specs, and filters — a single `.replace(src, tgt)` on the whole string covers all occurrences |
-
+| **`export_genie` / `import_genie` land in the wrong workspace** | Each MCP server is workspace-scoped. Set up two named MCP server entries (one per profile) in your IDE's MCP config instead of switching a single server's profile mid-session |
+| **MCP server doesn't pick up profile change** | The MCP process reads `DATABRICKS_CONFIG_PROFILE` once at startup — editing the config file requires an IDE reload to take effect |
+| **`import_genie` fails with JSON parse error** | The `serialized_space` string may contain multi-line SQL arrays with `\n` escape sequences; flatten SQL arrays to single-line strings before passing to avoid double-escaping issues |
 ## Related Skills
 
 - **[databricks-agent-bricks](../databricks-agent-bricks/SKILL.md)** - Use Genie Spaces as agents inside Supervisor Agents