Coverage for node / src / stigmem_node / cli / maintenance.py: 100%

1"""Maintenance and migration CLI handlers.""" 

2 

3from __future__ import annotations 

4 

5import argparse 

6import sys 

7 

8 

9def _cmd_decay_sweep(args: argparse.Namespace) -> int: 

10 from ..db import apply_migrations, db 

11 from ..lifecycle.decay import run_decay_sweep 

12 from ..settings import settings 

13 

14 db_path: str = args.db or settings.db_path 

15 apply_migrations(db_path=db_path) 

16 

17 if getattr(args, "all_tenants", False): 

18 # Multi-tenant operators: sweep every tenant so non-default tenants are 

19 # not silently skipped. Candidate selection inside run_decay_sweep is 

20 # tenant-scoped, so we iterate the distinct tenants and aggregate counts. 

21 with db() as conn: 

22 tenant_ids = [r["tenant_id"] for r in conn.execute( 

23 "SELECT DISTINCT tenant_id FROM facts" 

24 ).fetchall()] 

25 scanned = 0 

26 decayed = 0 

27 for tenant_id in tenant_ids: 

28 result = run_decay_sweep( 

29 ttl_seconds=args.ttl_seconds, 

30 min_confidence=args.min_confidence, 

31 scope=args.scope or None, 

32 dry_run=args.dry_run, 

33 tenant_id=tenant_id, 

34 ) 

35 scanned += result["scanned"] 

36 decayed += result["decayed"] 

37 else: 

38 result = run_decay_sweep( 

39 ttl_seconds=args.ttl_seconds, 

40 min_confidence=args.min_confidence, 

41 scope=args.scope or None, 

42 dry_run=args.dry_run, 

43 tenant_id=getattr(args, "tenant", "default") or "default", 

44 ) 

45 scanned = result["scanned"] 

46 decayed = result["decayed"] 

47 

48 if args.dry_run: 

49 print(f"[dry-run] {scanned} facts would be decayed", file=sys.stderr) 

50 else: 

51 print(f"{decayed} facts decayed ({scanned} scanned)", file=sys.stderr) 

52 return 0 

53 

54 

55def _cmd_migrate_normalize_entities(args: argparse.Namespace) -> int: 

56 from ..db import apply_migrations 

57 from ..migrate import normalize_entities_sweep 

58 from ..settings import settings 

59 

60 db_path: str = args.db or settings.db_path 

61 apply_migrations(db_path=db_path) 

62 

63 registered, already_present = normalize_entities_sweep(db_path, dry_run=args.dry_run) 

64 

65 if args.dry_run: 

66 print( 

67 f"[dry-run] {registered} aliases would be registered", 

68 file=sys.stderr, 

69 ) 

70 else: 

71 print( 

72 f"{registered} aliases registered, {already_present} already present", 

73 file=sys.stderr, 

74 ) 

75 return 0