mantatech-sdk 0.5b0.dev65__py3-none-any.whl

manta/cli/commands/sdk_results.py

@@ -0,0 +1,265 @@
+"""SDK results retrieval and management commands."""
+
+import json
+import signal
+import time
+
+from rich.table import Table
+
+from .base_handler import BaseSDKHandler
+
+
+class SDKResultsHandler(BaseSDKHandler):
+    """Handle results-related SDK commands."""
+
+    def add_subparsers(self, parent_parser):
+        """Add results command subparsers."""
+        self.parser = parent_parser  # Store parser reference
+        subparsers = parent_parser.add_subparsers(
+            dest="results_command", help="Results commands"
+        )
+
+        # Get results command
+        get_parser = subparsers.add_parser("get", help="Get results for a swarm")
+        get_parser.add_argument("swarm_id", help="Swarm ID to get results for")
+        get_parser.add_argument(
+            "--tag", default="latest", help="Result tag (default: latest)"
+        )
+
+        # Export results command
+        export_parser = subparsers.add_parser("export", help="Export results to file")
+        export_parser.add_argument("swarm_id", help="Swarm ID to export results for")
+        export_parser.add_argument(
+            "--tag", default="latest", help="Result tag (default: latest)"
+        )
+        export_parser.add_argument("--output", required=True, help="Output file path")
+
+        # Stream results command
+        stream_parser = subparsers.add_parser(
+            "stream", help="Stream results in real-time"
+        )
+        stream_parser.add_argument("swarm_id", help="Swarm ID to stream results for")
+        stream_parser.add_argument(
+            "--tag", default="latest", help="Result tag (default: latest)"
+        )
+
+        # List tags command
+        list_tags_parser = subparsers.add_parser(
+            "list-tags", help="List available result tags for a swarm"
+        )
+        list_tags_parser.add_argument("swarm_id", help="Swarm ID to list tags for")
+
+        # Add profile override to all subcommands
+        for parser in [get_parser, export_parser, stream_parser, list_tags_parser]:
+            parser.add_argument("--profile", help="Use specific profile")
+
+    def handle(self, args) -> int:
+        """Handle results commands."""
+        if not args.results_command:
+            if hasattr(self, "parser"):
+                self.parser.print_help()
+            return 0  # Return 0 for help display
+
+        if args.results_command == "get":
+            return self.get_results(args.swarm_id, getattr(args, "tag", "latest"))
+        elif args.results_command == "export":
+            return self.export_results(
+                args.swarm_id, getattr(args, "tag", "latest"), args.output
+            )
+        elif args.results_command == "stream":
+            return self.stream_results(args.swarm_id, getattr(args, "tag", "latest"))
+        elif args.results_command == "list-tags":
+            return self.list_tags(args.swarm_id)
+        else:
+            self.print_error(f"Unknown results command: {args.results_command}")
+            return 1
+
+    def get_results(self, swarm_id: str, tag: str = "latest") -> int:
+        """Get results for a swarm."""
+        api = self._get_user_api()
+        if not api:
+            return 1
+
+        try:
+
+            async def run():
+                return await api.get_results(swarm_id, tag)
+
+            results = self._run_with_progress(
+                run, f"Fetching results for swarm {swarm_id}..."
+            )
+
+            if results is None:
+                return 1
+
+            if not results:
+                self.print("No results found.")
+                return 0
+
+            # Display results
+            table = Table(title=f"Results for Swarm: {swarm_id} (Tag: {tag})")
+            table.add_column("Result ID", style="cyan")
+            table.add_column("Type", style="blue")
+            table.add_column("Size", style="green")
+            table.add_column("Created", style="yellow")
+
+            for result in results:
+                data = result.get("data", {})
+                data_size = len(json.dumps(data)) if data else 0
+
+                table.add_row(
+                    result.get("result_id", "Unknown"),
+                    result.get("result_type", "Unknown"),
+                    self._format_size(data_size),
+                    result.get("created_at", "Unknown"),
+                )
+
+            self.console.print(table)
+
+            # Show first few results as samples
+            self.print("\nSample Results:")
+            for i, result in enumerate(results[:3]):
+                self.print(f"Result {i + 1}:")
+                self.print(json.dumps(result.get("data", {}), indent=2))
+                if i < 2 and i < len(results) - 1:
+                    self.print("-" * 40)
+
+            return 0
+
+        except Exception as e:
+            self.print_error(f"Failed to get results: {e}")
+            return 1
+
+    def export_results(self, swarm_id: str, tag: str, output_file: str) -> int:
+        """Export results to a file."""
+        api = self._get_user_api()
+        if not api:
+            return 1
+
+        try:
+
+            async def run():
+                return await api.get_results(swarm_id, tag)
+
+            results = self._run_with_progress(run, "Fetching results for export...")
+
+            if results is None:
+                return 1
+
+            if not results:
+                self.print_error("No results found")
+                return 1
+
+            # Write to file
+            with open(output_file, "w") as f:
+                json.dump(results, f, indent=2)
+
+            self.print_success(f"Results exported to {output_file}")
+            return 0
+
+        except Exception as e:
+            self.print_error(f"Export failed: {e}")
+            return 1
+
+    def stream_results(self, swarm_id: str, tag: str = "latest") -> int:
+        """Stream results in real-time."""
+        api = self._get_user_api()
+        if not api:
+            return 1
+
+        self.print(f"Streaming results for swarm: {swarm_id} (Tag: {tag})")
+        self.print("Press Ctrl+C to stop streaming...")
+
+        try:
+            # Set up signal handler for graceful shutdown
+            shutdown_flag = False
+
+            def signal_handler(sig, frame):
+                nonlocal shutdown_flag
+                shutdown_flag = True
+
+            signal.signal(signal.SIGINT, signal_handler)
+
+            last_result_count = 0
+
+            while not shutdown_flag:
+                try:
+
+                    async def get_results():
+                        return await api.get_results(swarm_id, tag)
+
+                    results = self._run_async(get_results())
+
+                    # Show new results since last check
+                    if results and len(results) > last_result_count:
+                        new_results = results[last_result_count:]
+
+                        for result in new_results:
+                            timestamp = result.get("created_at", "unknown")
+                            data = result.get("data", {})
+                            self.print(f"[{timestamp}] {json.dumps(data, indent=2)}")
+
+                        last_result_count = len(results)
+
+                    time.sleep(5)  # Wait 5 seconds before next poll
+
+                except KeyboardInterrupt:
+                    self.print("\nStreaming stopped.")
+                    return 130  # KeyboardInterrupt exit code
+
+        except Exception as e:
+            self.print_error(f"Streaming failed: {e}")
+            return 1
+
+    def _format_size(self, size: int) -> str:
+        """Format byte size to human readable string."""
+        for unit in ["B", "KB", "MB", "GB"]:
+            if size < 1024.0:
+                return f"{size:.1f} {unit}"
+            size /= 1024.0
+        return f"{size:.1f} TB"
+
+    def list_tags(self, swarm_id: str) -> int:
+        """List available result tags for a swarm."""
+        api = self._get_user_api()
+        if not api:
+            return 1
+
+        try:
+
+            async def run():
+                return await api.list_result_tags(swarm_id)
+
+            tags_info = self._run_with_progress(
+                run, f"Fetching result tags for swarm {swarm_id}..."
+            )
+
+            if tags_info is None:
+                return 1
+
+            if not tags_info or not tags_info.get("tags"):
+                self.print("No result tags found for this swarm.")
+                return 0
+
+            # Display tags in a table
+            table = Table(title=f"Result Tags for Swarm: {swarm_id}")
+            table.add_column("Tag", style="cyan")
+            table.add_column("Count", style="green")
+            table.add_column("Total Size", style="yellow")
+            table.add_column("Last Updated", style="blue")
+
+            tags = tags_info.get("tags", {})
+            for tag, tag_data in tags.items():
+                count = tag_data.get("count", 0)
+                size = tag_data.get("total_size", 0)
+                last_updated = tag_data.get("last_updated", "Unknown")
+
+                table.add_row(tag, str(count), self._format_size(size), last_updated)
+
+            self.console.print(table)
+
+            return 0
+
+        except Exception as e:
+            self.print_error(f"Failed to list result tags: {e}")
+            return 1