#!/usr/bin/env python3
"""
Basic usage example for the Databricks MCP Server.
This script demonstrates how to use the DatabricksClient and UnityCatalogManager
to interact with Unity Catalog programmatically.
"""
import asyncio
import os
from dotenv import load_dotenv
from databricks_mcp_server import (
DatabricksClient,
UnityCatalogManager,
setup_logging,
validate_databricks_config
)
async def main():
"""Main example function."""
# Load environment variables
load_dotenv()
# Setup logging
setup_logging(level="INFO")
# Validate configuration
print("Validating Databricks configuration...")
validation = validate_databricks_config()
if not validation["valid"]:
print("❌ Configuration errors:")
for error in validation["errors"]:
print(f" - {error}")
return
if validation["warnings"]:
print("⚠️ Configuration warnings:")
for warning in validation["warnings"]:
print(f" - {warning}")
if validation["recommendations"]:
print("💡 Recommendations:")
for rec in validation["recommendations"]:
print(f" - {rec}")
print("✅ Configuration is valid!\n")
try:
# Initialize the Databricks client
print("Initializing Databricks client...")
client = await DatabricksClient.create()
print("✅ Client initialized successfully!\n")
# List catalogs
print("📚 Listing catalogs...")
catalogs = await client.list_catalogs()
print(f"Found {len(catalogs)} catalogs:")
for catalog in catalogs:
print(f" - {catalog.name}: {catalog.comment or 'No description'}")
print()
if catalogs:
# Get first catalog for examples
catalog_name = catalogs[0].name
assert catalog_name is not None, "Catalog name should not be None"
print(f"🔍 Exploring catalog '{catalog_name}'...")
# List schemas
schemas = await client.list_schemas(catalog_name)
print(f"Found {len(schemas)} schemas:")
for schema in schemas:
print(f" - {schema.name}: {schema.comment or 'No description'}")
print()
if schemas:
# Get first schema for examples
schema_name = schemas[0].name
assert schema_name is not None, "Schema name should not be None"
print(f"📋 Exploring schema '{catalog_name}.{schema_name}'...")
# List tables
tables = await client.list_tables(catalog_name, schema_name)
print(f"Found {len(tables)} tables:")
for table in tables:
table_type = table.table_type.value if table.table_type else "UNKNOWN"
print(f" - {table.name} ({table_type}): {table.comment or 'No description'}")
print()
if tables:
# Get detailed info for first table
table_name = tables[0].name
assert table_name is not None, "Table name should not be None"
print(f"🔍 Analyzing table '{catalog_name}.{schema_name}.{table_name}'...")
# Get table metadata
table_metadata = await client.describe_table(catalog_name, schema_name, table_name)
print(f"Table details:")
print(f" - Type: {table_metadata.table_type}")
print(f" - Owner: {table_metadata.owner or 'Unknown'}")
print(f" - Columns: {len(table_metadata.columns)}")
print(f" - Storage: {table_metadata.storage_location or 'Not specified'}")
if table_metadata.columns:
print(f" - Column details:")
for col in table_metadata.columns[:5]: # Show first 5 columns
nullable = "nullable" if col.get("nullable", True) else "not null"
print(f" * {col['name']} ({col['type']}, {nullable})")
if len(table_metadata.columns) > 5:
print(f" ... and {len(table_metadata.columns) - 5} more columns")
print()
# Sample data (if it's a table, not a view)
if table_metadata.table_type != "VIEW":
print(f"📊 Sampling data from '{table_name}'...")
try:
sample_result = await client.sample_table(catalog_name, schema_name, table_name, limit=3)
if sample_result.status == "SUCCESS":
print(f"Sample data ({sample_result.row_count} rows):")
for i, row in enumerate(sample_result.data):
print(f" Row {i+1}: {dict(list(row.items())[:3])}...") # Show first 3 columns
else:
print(f" ❌ Error sampling data: {sample_result.error}")
except Exception as e:
print(f" ⚠️ Could not sample data: {e}")
print()
# Initialize Unity Catalog manager for advanced operations
print("🚀 Demonstrating advanced features...")
uc_manager = UnityCatalogManager(client)
if catalogs:
# Get catalog summary
first_catalog_name = catalogs[0].name
assert first_catalog_name is not None, "First catalog name should not be None"
catalog_summary = await uc_manager.get_catalog_summary(first_catalog_name)
print(f"📈 Catalog '{catalog_summary.name}' summary:")
print(f" - Schemas: {catalog_summary.schema_count}")
print(f" - Tables: {catalog_summary.table_count}")
print(f" - Views: {catalog_summary.view_count}")
print()
# Search for tables
print("🔍 Searching for tables containing 'test' or 'sample'...")
search_results = await client.search_tables("test")
if search_results:
print(f"Found {len(search_results)} matching tables:")
for table in search_results[:3]: # Show first 3 results
print(f" - {table.catalog_name}.{table.schema_name}.{table.name}")
else:
print(" No tables found matching the search criteria")
print()
print("✅ Example completed successfully!")
except Exception as e:
print(f"❌ Error: {e}")
finally:
# Clean up
if 'client' in locals():
await client.close()
if __name__ == "__main__":
asyncio.run(main())