Coverage for session_buddy / tools / migration_tools.py: 40.00%

1#!/usr/bin/env python3 

2"""MCP tools for schema migration status and control.""" 

3 

4from __future__ import annotations 

5 

6import typing as t 

7from typing import TYPE_CHECKING 

8 

9from session_buddy.memory.migration import ( 

10 create_backup, 

11 get_migration_status, 

12 migrate_v1_to_v2, 

13 needs_migration, 

14) 

15 

16if TYPE_CHECKING: 

17 from fastmcp import FastMCP 

18 

19 

20def register_migration_tools(mcp: FastMCP) -> None: 

21 @mcp.tool() # type: ignore[no-untyped-call] 

22 async def migration_status() -> dict[str, t.Any]: 

23 """Check migration status and progress.""" 

24 return get_migration_status() 

25 

26 @mcp.tool() # type: ignore[no-untyped-call] 

27 async def trigger_migration( 

28 create_backup_first: bool = True, dry_run: bool = False 

29 ) -> dict[str, t.Any]: 

30 """Manually trigger migration (with preview).""" 

31 if create_backup_first and not dry_run: 

32 backup_path = create_backup() 

33 backup = str(backup_path) 

34 else: 

35 backup = None 

36 

37 result = migrate_v1_to_v2(dry_run=dry_run) 

38 return { 

39 "success": result.success, 

40 "error": result.error, 

41 "stats": result.stats or {}, 

42 "duration_seconds": result.duration_seconds, 

43 "backup": backup, 

44 "migration_needed": needs_migration() if not result.success else False, 

45 } 

46 

47 @mcp.tool() # type: ignore[no-untyped-call] 

48 async def rollback_migration(backup_path: str) -> dict[str, t.Any]: 

49 """Restore database from a previous backup file path.""" 

50 from pathlib import Path 

51 

52 from session_buddy.memory.migration import get_schema_version, restore_backup 

53 

54 restore_backup(Path(backup_path)) 

55 version = get_schema_version() 

56 status = get_migration_status() 

57 return {"restored_version": version, "status": status}