Yaizu Smart City MCP Server

#!/usr/bin/env python3 """ Databricks認証設定とセットアップスクリプト このスクリプトは、Databricks CLIとSDKの認証設定を支援します。 環境変数から認証情報を読み込み、複数のプロファイルを管理できます。 """ import os import sys import json import subprocess from pathlib import Path from typing import Optional, Dict, Any from dotenv import load_dotenv # 環境変数を読み込み load_dotenv() class DatabricksSetup: """Databricks環境セットアップユーティリティ""" def __init__(self): """初期化処理""" self.home_dir = Path.home() self.config_dir = self.home_dir / ".databricks" self.config_file = self.config_dir / "authentication.json" # 環境変数から設定を取得 self.host = os.getenv("DATABRICKS_HOST") self.token = os.getenv("DATABRICKS_TOKEN") self.profile = os.getenv("DATABRICKS_PROFILE", "DEFAULT") self.cluster_id = os.getenv("DATABRICKS_CLUSTER_ID") self.workspace_path = os.getenv("DATABRICKS_WORKSPACE_PATH") # Azure固有設定 self.azure_client_id = os.getenv("AZURE_CLIENT_ID") self.azure_client_secret = os.getenv("AZURE_CLIENT_SECRET") self.azure_tenant_id = os.getenv("AZURE_TENANT_ID") # AWS固有設定 self.aws_region = os.getenv("AWS_REGION") self.aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID") self.aws_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY") def check_cli_installed(self) -> bool: """Databricks CLIがインストールされているか確認""" try: result = subprocess.run( ["databricks", "-v"], capture_output=True, text=True, check=False ) if result.returncode == 0: print(f"✅ Databricks CLI確認: {result.stdout.strip()}") return True else: print("❌ Databricks CLIが見つかりません") return False except FileNotFoundError: print("❌ Databricks CLIがインストールされていません") return False def setup_auth_profile(self) -> bool: """認証プロファイルの設定""" if not self.host or not self.token: print("❌ 環境変数 DATABRICKS_HOST と DATABRICKS_TOKEN を設定してください") print(" .env.example を参考に .env ファイルを作成してください") return False print(f"\n🔐 認証プロファイル '{self.profile}' を設定中...") try: # databricks auth login コマンドで認証設定 # 注: このコマンドは対話的な認証フローを開始します cmd = [ "databricks", "auth", "token", "--host", self.host, "--token", self.token ] if self.profile != "DEFAULT": cmd.extend(["--profile", self.profile]) result = subprocess.run( cmd, capture_output=True, text=True, check=False ) if result.returncode == 0: print(f"✅ プロファイル '{self.profile}' の設定完了") return True else: print(f"❌ 認証設定エラー: {result.stderr}") return False except Exception as e: print(f"❌ 認証設定中にエラー: {e}") return False def test_connection(self) -> bool: """Databricks接続テスト""" print("\n🔍 接続テスト実行中...") try: # ワークスペースの現在のユーザー情報を取得 cmd = ["databricks", "current-user", "me"] if self.profile != "DEFAULT": cmd.extend(["--profile", self.profile]) result = subprocess.run( cmd, capture_output=True, text=True, check=False ) if result.returncode == 0: print("✅ 接続成功!") user_info = json.loads(result.stdout) print(f" ユーザー: {user_info.get('userName', 'Unknown')}") print(f" ホスト: {self.host}") return True else: print(f"❌ 接続失敗: {result.stderr}") return False except Exception as e: print(f"❌ 接続テスト中にエラー: {e}") return False def list_clusters(self) -> Optional[list]: """利用可能なクラスターのリスト取得""" print("\n📊 クラスター一覧取得中...") try: cmd = ["databricks", "clusters", "list"] if self.profile != "DEFAULT": cmd.extend(["--profile", self.profile]) result = subprocess.run( cmd, capture_output=True, text=True, check=False ) if result.returncode == 0: clusters = json.loads(result.stdout) print("✅ 利用可能なクラスター:") for cluster in clusters: state = cluster.get('state', 'UNKNOWN') name = cluster.get('cluster_name', 'Unknown') cluster_id = cluster.get('cluster_id', 'Unknown') print(f" - {name} (ID: {cluster_id}, State: {state})") return clusters else: print(f"❌ クラスター取得失敗: {result.stderr}") return None except Exception as e: print(f"❌ クラスター一覧取得中にエラー: {e}") return None def create_sample_notebook(self) -> bool: """サンプルノートブックの作成""" if not self.workspace_path: print("⚠️ DATABRICKS_WORKSPACE_PATH が設定されていないため、サンプル作成をスキップ") return True print(f"\n📝 サンプルノートブック作成中: {self.workspace_path}/hello_databricks.py") sample_code = """# Databricks notebook source # MAGIC %md # MAGIC # Hello Databricks! # MAGIC このノートブックは setup_databricks.py によって自動生成されました。 # COMMAND ---------- import datetime print(f"Hello from Databricks! 現在時刻: {datetime.datetime.now()}") # COMMAND ---------- # Spark DataFrame の簡単な例 from pyspark.sql import SparkSession spark = SparkSession.builder.appName("HelloDatabricks").getOrCreate() # サンプルデータの作成 data = [ ("Tokyo", "Japan", 37400000), ("Delhi", "India", 32900000), ("Shanghai", "China", 28500000), ("São Paulo", "Brazil", 22400000), ("Mexico City", "Mexico", 22300000) ] columns = ["city", "country", "population"] df = spark.createDataFrame(data, columns) # データの表示 df.show() df.printSchema() # COMMAND ---------- # 簡単な分析 print(f"総都市数: {df.count()}") print(f"総人口: {df.selectExpr('sum(population)').collect()[0][0]:,}") """ try: # ノートブックのアップロード import tempfile with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f: f.write(sample_code) temp_path = f.name cmd = [ "databricks", "workspace", "import", temp_path, f"{self.workspace_path}/hello_databricks.py", "--language", "PYTHON", "--overwrite" ] if self.profile != "DEFAULT": cmd.extend(["--profile", self.profile]) result = subprocess.run( cmd, capture_output=True, text=True, check=False ) # 一時ファイルを削除 os.unlink(temp_path) if result.returncode == 0: print("✅ サンプルノートブック作成完了") return True else: print(f"⚠️ ノートブック作成失敗（権限不足の可能性）: {result.stderr}") return False except Exception as e: print(f"⚠️ サンプルノートブック作成中にエラー: {e}") return False def print_next_steps(self): """次のステップを表示""" print("\n" + "="*50) print("📚 次のステップ:") print("="*50) print() print("1. 基本的なCLIコマンド:") print(" databricks current-user me # 現在のユーザー情報") print(" databricks clusters list # クラスター一覧") print(" databricks workspace ls / # ワークスペース内容表示") print(" databricks fs ls dbfs:/ # DBFS内容表示") print() print("2. ノートブック操作:") print(" databricks workspace export <path> # ノートブックのエクスポート") print(" databricks workspace import <file> # ノートブックのインポート") print() print("3. ジョブ実行:") print(" databricks jobs list # ジョブ一覧") print(" databricks jobs run-now --job-id <id> # ジョブ実行") print() print("4. Python SDKの使用:") print(" from databricks.sdk import WorkspaceClient") print(" w = WorkspaceClient()") print(" # SDKを使った操作が可能") print() print("詳細: https://docs.databricks.com/dev-tools/cli/index.html") def run(self) -> bool: """セットアップ実行""" print("🚀 Databricks セットアップ開始") print("="*50) # CLIインストール確認 if not self.check_cli_installed(): print("\n📦 Databricks CLIをインストールしてください:") print(" brew tap databricks/tap") print(" brew install databricks") return False # 認証設定 if not self.setup_auth_profile(): return False # 接続テスト if not self.test_connection(): return False # クラスター一覧表示 self.list_clusters() # サンプルノートブック作成 self.create_sample_notebook() # 次のステップを表示 self.print_next_steps() print("\n✨ Databricks セットアップ完了!") return True def main(): """メイン処理""" setup = DatabricksSetup() # .envファイルの確認 if not Path(".env").exists(): print("⚠️ .env ファイルが見つかりません") print(" .env.example を参考に .env ファイルを作成してください:") print(" cp .env.example .env") print(" その後、DATABRICKS_HOST と DATABRICKS_TOKEN を設定してください") sys.exit(1) # セットアップ実行 success = setup.run() sys.exit(0 if success else 1) if __name__ == "__main__": main()