Refactor create_or_update_genie function for improved clarity and functionality

- Streamline the logic for handling serialized_space, ensuring clearer separation of update and create operations.
- Enhance error handling for non-existent Genie spaces when updating by ID.
- Update documentation to reflect changes in the usage of ask_genie and ask_genie_followup tools, clarifying their purposes and examples.
- Revise examples in conversation and SKILL documentation to align with the new function structure and naming conventions.

Author: BlitzBricksterYY-db

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

@@ -107,82 +107,82 @@ 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, 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,
                 )
-            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", "")
         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", "")
-
-                # Add sample questions if provided
-                if sample_questions and space_id:
-                    manager.genie_add_sample_questions_batch(space_id, sample_questions)
+                # 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)
 
         response = {
             "space_id": space_id,

databricks-skills/databricks-genie/SKILL.md

@@ -33,7 +33,7 @@ Use this skill when:
 |------|---------|
 | `list_genie` | List all Genie Spaces accessible to you |
 | `create_or_update_genie` | Create or update a Genie Space (supports `serialized_space`) |
-| `get_genie` |  Get space details (by ID and support `include_serialized_space` parameter) or list all spaces (no ID) |
+| `get_genie` | Get Genie Space details (supports `include_serialized_space`) |
 | `delete_genie` | Delete a Genie Space |
 | `export_genie` | Export a Genie Space with full serialized configuration |
 | `import_genie` | Import / clone a Genie Space from a serialized payload |
@@ -42,7 +42,8 @@ Use this skill when:
 
 | Tool | Purpose |
 |------|---------|
-| `ask_genie` | Ask a question or follow-up (`conversation_id` optional) |
+| `ask_genie` | Ask a question to a Genie Space, get SQL + results |
+| `ask_genie_followup` | Ask follow-up question in existing conversation |
 
 ### Supporting Tools
 
@@ -113,28 +114,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:**
+
+```python
+# Call list_warehouses via the dev-scoped MCP server
+list_warehouses()  # note the warehouse_id for the DEV workspace
+```
 
-# 2. Remap the catalog name for the development environment
+**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
@@ -162,8 +227,8 @@ Before creating a Genie Space:
 ### Creating Tables
 
 Use these skills in sequence:
-1. `databricks-synthetic-data-gen` - Generate raw parquet files
-2. `databricks-spark-declarative-pipelines` - Create bronze/silver/gold tables
+1. `synthetic-data-generation` - Generate raw parquet files
+2. `spark-declarative-pipelines` - Create bronze/silver/gold tables
 
 ## Common Issues
 
@@ -176,10 +241,6 @@ 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 |
-
-## Related Skills
-
-- **[databricks-agent-bricks](../databricks-agent-bricks/SKILL.md)** - Use Genie Spaces as agents inside Supervisor Agents
-- **[databricks-synthetic-data-gen](../databricks-synthetic-data-gen/SKILL.md)** - Generate raw parquet data to populate tables for Genie
-- **[databricks-spark-declarative-pipelines](../databricks-spark-declarative-pipelines/SKILL.md)** - Build bronze/silver/gold tables consumed by Genie Spaces
-- **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - Manage the catalogs, schemas, and tables Genie queries
+| **`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 |

databricks-skills/databricks-genie/conversation.md

@@ -31,7 +31,8 @@ The `ask_genie` tool allows you to programmatically send questions to a Genie Sp
 
 | Tool | Purpose |
 |------|---------|
-| `ask_genie` | Ask a question or follow-up (`conversation_id` optional) |
+| `ask_genie` | Ask a question, start new conversation |
+| `ask_genie_followup` | Ask follow-up in existing conversation |
 
 ## Basic Usage
 
@@ -70,10 +71,10 @@ result = ask_genie(
 )
 
 # Follow-up (uses context from first question)
-ask_genie(
+ask_genie_followup(
     space_id="01abc123...",
-    question="Break that down by region",
-    conversation_id=result["conversation_id"]
+    conversation_id=result["conversation_id"],
+    question="Break that down by region"
 )
 ```
 
@@ -162,9 +163,9 @@ User: "Use my analytics Genie to explore sales trends"
 Claude:
 1. ask_genie(space_id, "What were total sales by month this year?")
 2. User: "Which month had the highest growth?"
-3. ask_genie(space_id, "Which month had the highest growth?", conversation_id=conv_id)
+3. ask_genie_followup(space_id, conv_id, "Which month had the highest growth?")
 4. User: "What products drove that growth?"
-5. ask_genie(space_id, "What products drove that growth?", conversation_id=conv_id)
+5. ask_genie_followup(space_id, conv_id, "What products drove that growth?")
 ```
 
 ## Best Practices
@@ -180,8 +181,8 @@ result2 = ask_genie(space_id, "How many employees do we have?")  # New conversat
 
 # Good: Follow-up for related question
 result1 = ask_genie(space_id, "What were sales last month?")
-result2 = ask_genie(space_id, "Break that down by product",
-                    conversation_id=result1["conversation_id"])  # Related follow-up
+result2 = ask_genie_followup(space_id, result1["conversation_id"],
+                              "Break that down by product")  # Related follow-up
 ```
 
 ### Handle Clarification Requests

databricks-skills/databricks-genie/spaces.md

@@ -297,10 +297,10 @@ create_or_update_genie(
 
 ## Example End-to-End Workflow
 
-1. **Generate synthetic data** using `databricks-synthetic-data-gen` skill:
+1. **Generate synthetic data** using `synthetic-data-generation` skill:
    - Creates parquet files in `/Volumes/catalog/schema/raw_data/`
 
-2. **Create tables** using `databricks-spark-declarative-pipelines` skill:
+2. **Create tables** using `spark-declarative-pipelines` skill:
    - Creates `catalog.schema.bronze_*` → `catalog.schema.silver_*` → `catalog.schema.gold_*`
 
 3. **Inspect the tables**: