#!/usr/bin/env python3
"""
Example usage of Databricks MCP Server Prompts
This example demonstrates how to use the various prompts available in the
Databricks MCP server for common data workflows.
"""
import asyncio
import json
from databricks_mcp_server.prompts import PromptsManager
async def main():
"""Demonstrate prompt usage."""
prompts_manager = PromptsManager()
print("=== Available Databricks Prompts ===")
prompts = prompts_manager.list_prompts()
for prompt in prompts:
print(f"\nš {prompt.name}")
print(f" Description: {prompt.description}")
if prompt.arguments:
print(" Arguments:")
for arg in prompt.arguments:
required = "Required" if arg.required else "Optional"
print(f" - {arg.name} ({required}): {arg.description}")
print("\n" + "="*60)
print("=== Example Prompt Usage ===")
# Example 1: Dataset exploration prompt
print("\n1. Dataset Exploration Prompt:")
try:
explore_result = prompts_manager.get_prompt(
"explore-dataset",
{
"catalog_name": "my_catalog",
"schema_name": "sales_data",
"table_name": "customer_orders",
"analysis_focus": "data quality and business insights"
}
)
print(f"Description: {explore_result.description}")
print(f"Generated prompt:\n{explore_result.messages[0].content.text[:500]}...")
except Exception as e:
print(f"Error: {e}")
# Example 2: Query optimization prompt
print("\n2. Query Optimization Prompt:")
try:
optimize_result = prompts_manager.get_prompt(
"optimize-query",
{
"query": """
SELECT c.customer_id, c.name, COUNT(o.order_id) as order_count
FROM customers c
LEFT JOIN orders o ON c.customer_id = o.customer_id
WHERE c.created_date >= '2024-01-01'
GROUP BY c.customer_id, c.name
ORDER BY order_count DESC
""",
"performance_issues": "Query takes 5+ minutes to execute",
"execution_context": "Used in daily executive dashboard"
}
)
print(f"Description: {optimize_result.description}")
print(f"Generated prompt:\n{optimize_result.messages[0].content.text[:500]}...")
except Exception as e:
print(f"Error: {e}")
# Example 3: Data quality checks prompt
print("\n3. Data Quality Checks Prompt:")
try:
quality_result = prompts_manager.get_prompt(
"design-data-quality-checks",
{
"catalog_name": "my_catalog",
"schema_name": "sales_data",
"table_name": "customer_orders",
"business_rules": "Orders must have positive amounts, valid customer IDs, and order dates within business hours",
"quality_dimensions": "completeness, accuracy, consistency"
}
)
print(f"Description: {quality_result.description}")
print(f"Generated prompt:\n{quality_result.messages[0].content.text[:500]}...")
except Exception as e:
print(f"Error: {e}")
# Example 4: Documentation generation prompt
print("\n4. Documentation Generation Prompt:")
try:
doc_result = prompts_manager.get_prompt(
"generate-documentation",
{
"catalog_name": "my_catalog",
"schema_name": "sales_data",
"table_name": "customer_orders",
"audience": "business analysts and data scientists",
"doc_style": "business-friendly"
}
)
print(f"Description: {doc_result.description}")
print(f"Generated prompt:\n{doc_result.messages[0].content.text[:500]}...")
except Exception as e:
print(f"Error: {e}")
# Example 5: Troubleshooting prompt
print("\n5. Troubleshooting Prompt:")
try:
trouble_result = prompts_manager.get_prompt(
"troubleshoot-data-issue",
{
"issue_description": "Customer order counts don't match between staging and production tables",
"affected_table": "my_catalog.sales_data.customer_orders",
"observed_symptoms": "Production shows 10% fewer orders than staging for the same date range",
"timeframe": "Issue noticed on 2024-01-15, affects data from 2024-01-10 onwards"
}
)
print(f"Description: {trouble_result.description}")
print(f"Generated prompt:\n{trouble_result.messages[0].content.text[:500]}...")
except Exception as e:
print(f"Error: {e}")
print("\n" + "="*60)
print("=== Integration with MCP Client ===")
print("""
To use these prompts with an MCP client:
1. List available prompts:
Client sends: { "method": "prompts/list" }
Server responds with list of available prompts
2. Get a specific prompt:
Client sends: {
"method": "prompts/get",
"params": {
"name": "explore-dataset",
"arguments": {
"catalog_name": "my_catalog",
"schema_name": "sales_data",
"table_name": "customer_orders"
}
}
}
Server responds with formatted prompt messages
3. Client can then use the prompt messages to guide LLM interaction
and utilize Databricks tools to fulfill the prompt instructions.
""")
if __name__ == "__main__":
asyncio.run(main())