Admin/XMover: Implement suggestions by CodeRabbit

Author: amotl

cratedb_toolkit/admin/xmover/analysis/shard.py

@@ -38,7 +38,7 @@ def __init__(self, client: CrateDBClient):
         self.shards: List[ShardInfo] = []
 
         # Initialize session-based caches for performance.
-        self._zone_conflict_cache: Dict[Tuple[str, int, str], Union[str, None]] = {}
+        self._zone_conflict_cache: Dict[Tuple[str, str, int, str], Union[str, None]] = {}
         self._node_lookup_cache: Dict[str, Union[NodeInfo, None]] = {}
         self._target_nodes_cache: Dict[Tuple[float, frozenset[Any], float, float], List[NodeInfo]] = {}
         self._cache_hits = 0
@@ -183,8 +183,6 @@ def find_nodes_with_capacity(
             free_space_gb = node.available_space_gb
             if free_space_gb >= (required_space_gb + min_free_space_gb):
                 available_nodes.append(node)
-            else:
-                continue
 
         # Sort by available space (most space first) - prioritize nodes with more free space
         available_nodes.sort(key=lambda n: n.available_space_gb, reverse=True)
@@ -206,7 +204,7 @@ def generate_rebalancing_recommendations(
         # Get moveable shards (only healthy ones for actual operations)
         moveable_shards = self.find_moveable_shards(constraints.min_size, constraints.max_size, constraints.table_name)
 
-        print(
+        logger.info(
             f"Analyzing {len(moveable_shards)} candidate shards "
             f"in size range {constraints.min_size}-{constraints.max_size}GB..."
         )
@@ -239,12 +237,11 @@ def generate_rebalancing_recommendations(
         # Optimize processing: if filtering by source node, only process those shards
         if constraints.source_node:
             processing_shards = [s for s in moveable_shards if s.node_name == constraints.source_node]
-            print(f"Focusing on {len(processing_shards)} shards from node {constraints.source_node}")
+            logger.info(f"Focusing on {len(processing_shards)} shards from node {constraints.source_node}")
         else:
             processing_shards = moveable_shards
 
         # Generate move recommendations
-        safe_recommendations = 0  # noqa: F841
         total_evaluated = 0
 
         for i, shard in enumerate(processing_shards):
@@ -368,12 +365,12 @@ def generate_rebalancing_recommendations(
 
         if len(processing_shards) > 100:
             print()  # New line after progress dots
-        print(f"Generated {len(recommendations)} move recommendations (evaluated {total_evaluated} shards)")
-        print(f"Performance: {self.get_cache_stats()}")
+        logger.info(f"Generated {len(recommendations)} move recommendations (evaluated {total_evaluated} shards)")
+        logger.info(f"Performance: {self.get_cache_stats()}")
         return recommendations
 
     def validate_move_safety(
-        self, recommendation: ShardRelocationResponse, max_disk_usage_percent: float = 90.0
+        self, recommendation: ShardRelocationResponse, max_disk_usage_percent: float = 90.0, buffer_gb: float = 50.0
     ) -> Tuple[bool, str]:
         """Validate that a move recommendation is safe to execute"""
         # Find target node (with caching)
@@ -388,7 +385,7 @@ def validate_move_safety(
             return False, zone_conflict
 
         # Check available space
-        required_space_gb = recommendation.size_gb + 50  # 50GB buffer
+        required_space_gb = recommendation.size_gb + buffer_gb
         if target_node.available_space_gb < required_space_gb:
             return (
                 False,
@@ -423,7 +420,7 @@ def _check_zone_conflict_cached(self, recommendation: ShardRelocationResponse) -
         """Check zone conflicts with caching"""
         # Create cache key: table, shard, target zone
         target_zone = self._get_node_zone(recommendation.to_node)
-        cache_key = (recommendation.table_name, recommendation.shard_id, target_zone)
+        cache_key = (recommendation.schema_name, recommendation.table_name, recommendation.shard_id, target_zone)
 
         if cache_key in self._zone_conflict_cache:
             self._cache_hits += 1
@@ -813,11 +810,14 @@ def plan_node_decommission(self, node_name: str, min_free_space_gb: float = 100.
         # Determine feasibility
         feasible = len(infeasible_moves) == 0
 
+        # Safety margin for cluster capacity after decommission
+        capacity_safety_margin = 1.2  # 20 % buffer
+
         # Add capacity warnings
         if feasible:
-            # Check if remaining cluster capacity is sufficient after decommission
+            # Check if the remaining cluster capacity is sufficient after decommission
             remaining_capacity = sum(n.available_space_gb for n in self.nodes if n.name != node_name)
-            if remaining_capacity < total_size_gb * 1.2:  # 20% safety margin
+            if remaining_capacity < total_size_gb * capacity_safety_margin:
                 warnings.append(
                     f"Low remaining capacity after decommission. "
                     f"Only {remaining_capacity:.1f}GB available for {total_size_gb:.1f}GB of data"
@@ -833,7 +833,7 @@ def plan_node_decommission(self, node_name: str, min_free_space_gb: float = 100.
             "recommendations": move_plan,
             "infeasible_moves": infeasible_moves,
             "warnings": warnings,
-            "estimated_time_hours": len(move_plan) * 0.1,  # Rough estimate: 6 minutes per move
+            "estimated_time_hours": len(move_plan) * 0.1,  # Rough estimate: 0.1 hours (6 minutes) per move
             "message": "Decommission plan generated" if feasible else "Decommission not currently feasible",
         }
 

cratedb_toolkit/admin/xmover/analysis/table.py

@@ -14,6 +14,7 @@
 from rich.console import Console
 from rich.table import Table
 
+from cratedb_toolkit.admin.xmover.model import NodeInfo
 from cratedb_toolkit.admin.xmover.util.database import CrateDBClient
 
 logger = logging.getLogger(__name__)
@@ -97,6 +98,9 @@ def find_table_by_name(self, table_name: str) -> Optional[str]:
 
             try:
                 choice = input("\nSelect table (enter number): ").strip()
+                if not choice:
+                    rprint("[yellow]No selection made[/yellow]")
+                    return None
                 idx = int(choice) - 1
                 if 0 <= idx < len(rows):
                     schema, table = rows[idx]
@@ -292,14 +296,9 @@ def format_table_health_report(self, table_dist: TableDistribution) -> None:
             zone_distribution = {}
             for node_name, node_data in table_dist.node_distributions.items():
                 # Try to get zone info for each node
-                node_info = next((n for n in all_nodes_info if n.name == node_name), None)
-                if (
-                    node_info
-                    and hasattr(node_info, "attributes")
-                    and node_info.attributes
-                    and "zone" in node_info.attributes
-                ):
-                    zone = node_info.attributes["zone"]
+                node_info: Optional[NodeInfo] = next((n for n in all_nodes_info if n.name == node_name), None)
+                if node_info and node_info.zone:
+                    zone = node_info.zone
                     if zone not in zone_distribution:
                         zone_distribution[zone] = {"nodes": 0, "shards": 0, "size": 0}
                     zone_distribution[zone]["nodes"] += 1

cratedb_toolkit/admin/xmover/analysis/zone.py

@@ -135,7 +135,8 @@ def distribution_conflicts(self, shard_details: bool = False, table: Optional[st
                         health_indicator = "✓" if shard_copy.routing_state == "STARTED" else "⚠"
                         console.print(
                             f"    {health_indicator} {shard_copy.shard_type} "
-                            f"on {shard_copy.node_name} ({shard_copy.zone}) - {shard_copy.routing_state}"
+                            f"on {shard_copy.node_name} ({shard_copy.zone}) - "
+                            f"{shard_copy.state}/{shard_copy.routing_state}"
                         )
 
             console.print(analysis_table)

cratedb_toolkit/admin/xmover/cli.py

@@ -4,7 +4,6 @@
 Command Line Interface.
 """
 
-import sys
 import time
 from typing import Optional
 
@@ -46,11 +45,11 @@ def main(ctx):
         if not client.test_connection():
             console.print("[red]Error: Could not connect to CrateDB[/red]")
             console.print("Please check your CRATE_CONNECTION_STRING in .env file")
-            sys.exit(1)
+            raise click.Abort()
         ctx.obj["client"] = client
     except Exception as e:
         console.print(f"[red]Error connecting to CrateDB: {e}[/red]")
-        sys.exit(1)
+        raise click.Abort() from e
 
 
 @main.command()
@@ -170,11 +169,11 @@ def test_connection(ctx, connection_string: Optional[str]):
                 console.print(f"  • {node.name} (zone: {node.zone})")
         else:
             console.print("[red]✗ Connection failed[/red]")
-            sys.exit(1)
+            raise click.Abort()
 
     except Exception as e:
         console.print(f"[red]✗ Connection error: {e}[/red]")
-        sys.exit(1)
+        raise click.Abort() from e
 
 
 @main.command()
@@ -525,13 +524,14 @@ def monitor_recovery(
         xmover monitor-recovery --watch                # Continuous monitoring
         xmover monitor-recovery --recovery-type PEER  # Only PEER recoveries
     """
+    effective_recovery_type = None if recovery_type == "all" else recovery_type
     recovery_monitor = RecoveryMonitor(
         client=ctx.obj["client"],
         options=RecoveryOptions(
             table=table,
             node=node,
             refresh_interval=refresh_interval,
-            recovery_type=recovery_type,
+            recovery_type=effective_recovery_type,
             include_transitioning=include_transitioning,
         ),
     )

cratedb_toolkit/admin/xmover/model.py

@@ -1,4 +1,3 @@
-import dataclasses
 from dataclasses import dataclass
 from typing import Dict, Optional
 
@@ -149,6 +148,12 @@ def safety_score(self) -> float:
         if "rebalancing" in self.reason.lower():
             score += 0.2
 
+        # Consider shard size - smaller shards are safer to move
+        if self.size_gb < 10:
+            score += 0.1
+        elif self.size_gb > 100:
+            score -= 0.2
+
         # Ensure score stays in valid range
         return max(0.0, min(1.0, score))
 
@@ -165,15 +170,15 @@ class DistributionStats:
     node_balance_score: float  # 0-100, higher is better
 
 
-@dataclasses.dataclass
+@dataclass
 class SizeCriteria:
     min_size: float = 40.0
     max_size: float = 60.0
     table_name: Optional[str] = None
     source_node: Optional[str] = None
 
 
-@dataclasses.dataclass
+@dataclass
 class ShardRelocationConstraints:
     min_size: float = SizeCriteria().min_size
     max_size: float = SizeCriteria().max_size

cratedb_toolkit/admin/xmover/operational/candidates.py

@@ -14,7 +14,7 @@ class CandidateFinder:
     def __init__(self, analyzer: ShardAnalyzer):
         self.analyzer = analyzer
 
-    def movement_candidates(self, criteria: SizeCriteria, limit: int):
+    def movement_candidates(self, criteria: SizeCriteria, limit: int) -> int:
         """
         Find shard candidates for movement based on size criteria
 
@@ -23,7 +23,7 @@ def movement_candidates(self, criteria: SizeCriteria, limit: int):
         """
 
         console.print(
-            Panel.fit(f"[bold blue]Finding Moveable Shards ({criteria.min_size}-{criteria.max_size}GB)[/bold blue]")
+            Panel.fit(f"[bold blue]Finding Movable Shards ({criteria.min_size}-{criteria.max_size}GB)[/bold blue]")
         )
 
         if criteria.source_node:
@@ -45,7 +45,7 @@ def movement_candidates(self, criteria: SizeCriteria, limit: int):
                 console.print("[dim]Tip: Try different size ranges or remove --node filter to see all candidates[/dim]")
             else:
                 console.print("[yellow]No moveable shards found in the specified size range.[/yellow]")
-            return
+            return 0
 
         # Show limited results
         shown_candidates = candidates[:limit]
@@ -82,3 +82,5 @@ def movement_candidates(self, criteria: SizeCriteria, limit: int):
 
         if len(candidates) > limit:
             console.print(f"\n[dim]... and {len(candidates) - limit} more candidates[/dim]")
+
+        return len(candidates)

cratedb_toolkit/admin/xmover/operational/monitor.py

@@ -37,7 +37,7 @@ def get_cluster_recovery_status(self) -> List[RecoveryInfo]:
         )
 
         # Apply recovery type filter
-        if self.options.recovery_type is not None:
+        if self.options.recovery_type is not None and self.options.recovery_type.lower() != "all":
             recoveries = [r for r in recoveries if r.recovery_type.upper() == self.options.recovery_type.upper()]
 
         return recoveries
@@ -178,7 +178,6 @@ def start(self, watch: bool, debug: bool = False):
 
                     # Track previous state for change detection
                     previous_recoveries: Dict[str, Dict[str, Any]] = {}
-                    previous_timestamp = None
                     first_run = True
 
                     while True:
@@ -307,7 +306,6 @@ def start(self, watch: bool, debug: bool = False):
                                 elif active_count > 0:
                                     console.print(f"{current_time} | {status} (no changes)")
 
-                        previous_timestamp = current_time  # noqa: F841
                         first_run = False
                         time.sleep(self.options.refresh_interval)
 

cratedb_toolkit/admin/xmover/operational/recommend.py

@@ -123,8 +123,8 @@ def validate(self, request: ShardRelocationRequest):
             console.print()
             console.print("[dim]# Monitor shard health after execution[/dim]")
             console.print(
-                "[dim]# Check with: SELECT * FROM sys.shards "
-                "WHERE table_name = '{table_name}' AND id = {shard_id};[/dim]"
+                "[dim]# Check with: SELECT * FROM sys.shards "  # noqa: S608
+                f"WHERE table_name = '{table_name}' AND id = {request.shard_id};[/dim]"
             )
         else:
             console.print("[red]✗ VALIDATION FAILED - Move not safe[/red]")
@@ -323,7 +323,7 @@ def execute(
                         rec, max_disk_usage_percent=constraints.max_disk_usage
                     )
                     if not is_safe:
-                        if "Zone conflict" in safety_msg:
+                        if "zone conflict" in safety_msg.lower():
                             zone_conflicts += 1
                             console.print(f"-- Move {i}: SKIPPED - {safety_msg}")
                             console.print(
@@ -340,7 +340,7 @@ def execute(
 
             # Auto-execution if requested
             if auto_execute:
-                self._execute_recommendations_safely(recommendations, validate)
+                self._execute_recommendations_safely(constraints, recommendations, validate)
 
         if validate and safe_moves < len(recommendations):
             if zone_conflicts > 0:
@@ -352,14 +352,16 @@ def execute(
                 f"[yellow]Warning: Only {safe_moves} of {len(recommendations)} moves passed safety validation[/yellow]"
             )
 
-    def _execute_recommendations_safely(self, recommendations, validate: bool):
+    def _execute_recommendations_safely(self, constraints, recommendations, validate: bool):
         """Execute recommendations with extensive safety measures"""
 
         # Filter to only safe recommendations
         safe_recommendations = []
         if validate:
             for rec in recommendations:
-                is_safe, safety_msg = self.analyzer.validate_move_safety(rec, max_disk_usage_percent=95.0)
+                is_safe, safety_msg = self.analyzer.validate_move_safety(
+                    rec, max_disk_usage_percent=constraints.max_disk_usage
+                )
                 if is_safe:
                     safe_recommendations.append(rec)
         else:
@@ -423,7 +425,8 @@ def _execute_recommendations_safely(self, recommendations, validate: bool):
                 # Execute the SQL command
                 result = self.client.execute_query(sql_command)
 
-                if result.get("rowcount", 0) >= 0:  # Success indicator for ALTER statements
+                # ALTER TABLE REROUTE commands don't return rowcount, check for no error instead.
+                if "error" not in result:
                     console.print("    [green]✅ SUCCESS[/green] - Move initiated")
                     successful_moves += 1
 
@@ -482,7 +485,8 @@ def _wait_for_recovery_capacity(self, max_concurrent_recoveries: int = 5):
         while True:
             # Check active recoveries (including transitioning)
             recoveries = recovery_monitor.get_cluster_recovery_status()
-            active_count = len([r for r in recoveries if r.overall_progress < 100.0 or r.stage != "DONE"])
+            # Count recoveries that are actively running (not completed)
+            active_count = len([r for r in recoveries if r.overall_progress < 100.0])
             status = f"{active_count}/{max_concurrent_recoveries}"
             if active_count < max_concurrent_recoveries:
                 if wait_time > 0: