pltr-cli 0.1.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

pltr/commands/sql.py

@@ -0,0 +1,358 @@
+"""
+SQL commands for the pltr CLI.
+Provides commands for executing SQL queries against Foundry datasets.
+"""
+
+from pathlib import Path
+from typing import Optional
+
+import typer
+from rich.console import Console
+
+from ..services.sql import SqlService
+from ..utils.formatting import OutputFormatter
+from ..utils.progress import SpinnerProgressTracker
+from ..utils.completion import (
+    complete_profile,
+    complete_output_format,
+    complete_sql_query,
+)
+
+app = typer.Typer(name="sql", help="Execute SQL queries against Foundry datasets")
+
+
+@app.command("execute")
+def execute_query(
+    query: str = typer.Argument(
+        ..., help="SQL query to execute", autocompletion=complete_sql_query
+    ),
+    profile: Optional[str] = typer.Option(
+        None, "--profile", help="Auth profile to use", autocompletion=complete_profile
+    ),
+    output_format: str = typer.Option(
+        "table",
+        "--format",
+        help="Output format (table, json, csv)",
+        autocompletion=complete_output_format,
+    ),
+    output_file: Optional[Path] = typer.Option(
+        None, "--output", help="Save results to file"
+    ),
+    timeout: int = typer.Option(300, "--timeout", help="Query timeout in seconds"),
+    fallback_branches: Optional[str] = typer.Option(
+        None, "--fallback-branches", help="Comma-separated list of fallback branch IDs"
+    ),
+) -> None:
+    """Execute a SQL query and display results."""
+    console = Console()
+    formatter = OutputFormatter()
+
+    try:
+        # Parse fallback branches if provided
+        fallback_branch_ids = (
+            fallback_branches.split(",") if fallback_branches else None
+        )
+
+        # Create service and execute query
+        service = SqlService(profile=profile)
+
+        with SpinnerProgressTracker().track_spinner("Executing SQL query..."):
+            result = service.execute_query(
+                query=query,
+                fallback_branch_ids=fallback_branch_ids,
+                timeout=timeout,
+                format="table" if output_format in ["table", "csv"] else "json",
+            )
+
+        # Extract results
+        query_results = result.get("results", {})
+
+        # Display results
+        if output_file:
+            formatter.save_to_file(query_results, output_file, output_format)
+            console.print(f"[green]Results saved to {output_file}[/green]")
+        else:
+            formatter.display(query_results, output_format)
+
+        # Show query metadata
+        if "query_id" in result:
+            console.print(f"\n[dim]Query ID: {result['query_id']}[/dim]")
+
+    except Exception as e:
+        formatter.print_error(f"Failed to execute query: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("submit")
+def submit_query(
+    query: str = typer.Argument(..., help="SQL query to submit"),
+    profile: Optional[str] = typer.Option(
+        None, "--profile", help="Auth profile to use"
+    ),
+    fallback_branches: Optional[str] = typer.Option(
+        None, "--fallback-branches", help="Comma-separated list of fallback branch IDs"
+    ),
+) -> None:
+    """Submit a SQL query without waiting for completion."""
+    console = Console()
+    formatter = OutputFormatter()
+
+    try:
+        # Parse fallback branches if provided
+        fallback_branch_ids = (
+            fallback_branches.split(",") if fallback_branches else None
+        )
+
+        # Create service and submit query
+        service = SqlService(profile=profile)
+
+        with SpinnerProgressTracker().track_spinner("Submitting SQL query..."):
+            result = service.submit_query(
+                query=query, fallback_branch_ids=fallback_branch_ids
+            )
+
+        console.print("[green]Query submitted successfully[/green]")
+        console.print(f"Query ID: [bold]{result.get('query_id', 'N/A')}[/bold]")
+        console.print(f"Status: [yellow]{result.get('status', 'unknown')}[/yellow]")
+
+        if result.get("status") == "succeeded":
+            console.print("[green]Query completed immediately[/green]")
+        elif result.get("status") == "running":
+            console.print(
+                "Use [bold]pltr sql status <query-id>[/bold] to check progress"
+            )
+            console.print(
+                "Use [bold]pltr sql results <query-id>[/bold] to get results when completed"
+            )
+
+    except Exception as e:
+        formatter.print_error(f"Failed to submit query: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("status")
+def get_query_status(
+    query_id: str = typer.Argument(..., help="Query ID to check"),
+    profile: Optional[str] = typer.Option(
+        None, "--profile", help="Auth profile to use"
+    ),
+) -> None:
+    """Get the status of a submitted query."""
+    console = Console()
+    formatter = OutputFormatter()
+
+    try:
+        service = SqlService(profile=profile)
+
+        with SpinnerProgressTracker().track_spinner("Checking query status..."):
+            result = service.get_query_status(query_id)
+
+        console.print(f"Query ID: [bold]{query_id}[/bold]")
+
+        status = result.get("status", "unknown")
+        if status == "running":
+            console.print(f"Status: [yellow]{status}[/yellow]")
+            console.print("Query is still executing...")
+        elif status == "succeeded":
+            console.print(f"Status: [green]{status}[/green]")
+            console.print("Use [bold]pltr sql results <query-id>[/bold] to get results")
+        elif status == "failed":
+            console.print(f"Status: [red]{status}[/red]")
+            error_msg = result.get("error_message", "Unknown error")
+            console.print(f"Error: {error_msg}")
+        elif status == "canceled":
+            console.print(f"Status: [red]{status}[/red]")
+            console.print("Query was canceled")
+        else:
+            console.print(f"Status: [dim]{status}[/dim]")
+
+    except Exception as e:
+        formatter.print_error(f"Failed to get query status: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("results")
+def get_query_results(
+    query_id: str = typer.Argument(..., help="Query ID to get results for"),
+    profile: Optional[str] = typer.Option(
+        None, "--profile", help="Auth profile to use"
+    ),
+    output_format: str = typer.Option(
+        "table", "--format", help="Output format (table, json, csv)"
+    ),
+    output_file: Optional[Path] = typer.Option(
+        None, "--output", help="Save results to file"
+    ),
+) -> None:
+    """Get the results of a completed query."""
+    console = Console()
+    formatter = OutputFormatter()
+
+    try:
+        service = SqlService(profile=profile)
+
+        with SpinnerProgressTracker().track_spinner("Retrieving query results..."):
+            result = service.get_query_results(
+                query_id,
+                format="table" if output_format in ["table", "csv"] else "json",
+            )
+
+        # Display or save results
+        if output_file:
+            formatter.save_to_file(result, output_file, output_format)
+            console.print(f"[green]Results saved to {output_file}[/green]")
+        else:
+            formatter.display(result, output_format)
+
+        console.print(f"\n[dim]Query ID: {query_id}[/dim]")
+
+    except Exception as e:
+        formatter.print_error(f"Failed to get query results: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("cancel")
+def cancel_query(
+    query_id: str = typer.Argument(..., help="Query ID to cancel"),
+    profile: Optional[str] = typer.Option(
+        None, "--profile", help="Auth profile to use"
+    ),
+) -> None:
+    """Cancel a running query."""
+    console = Console()
+    formatter = OutputFormatter()
+
+    try:
+        service = SqlService(profile=profile)
+
+        with SpinnerProgressTracker().track_spinner("Canceling query..."):
+            result = service.cancel_query(query_id)
+
+        console.print(f"Query ID: [bold]{query_id}[/bold]")
+
+        status = result.get("status", "unknown")
+        if status == "canceled":
+            console.print(f"Status: [red]{status}[/red]")
+            console.print("Query has been canceled successfully")
+        else:
+            console.print(f"Status: [yellow]{status}[/yellow]")
+            console.print("Query may have already completed or was not running")
+
+    except Exception as e:
+        formatter.print_error(f"Failed to cancel query: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("export")
+def export_query_results(
+    query: str = typer.Argument(..., help="SQL query to execute and export"),
+    output_file: Path = typer.Argument(..., help="Output file path"),
+    profile: Optional[str] = typer.Option(
+        None, "--profile", help="Auth profile to use"
+    ),
+    output_format: Optional[str] = typer.Option(
+        None,
+        "--format",
+        help="Output format (auto-detected from file extension if not specified)",
+    ),
+    timeout: int = typer.Option(300, "--timeout", help="Query timeout in seconds"),
+    fallback_branches: Optional[str] = typer.Option(
+        None, "--fallback-branches", help="Comma-separated list of fallback branch IDs"
+    ),
+) -> None:
+    """Execute a SQL query and export results to a file."""
+    console = Console()
+    formatter = OutputFormatter()
+
+    try:
+        # Auto-detect format from file extension if not specified
+        if output_format is None:
+            ext = output_file.suffix.lower()
+            if ext == ".json":
+                output_format = "json"
+            elif ext == ".csv":
+                output_format = "csv"
+            else:
+                output_format = "table"  # Default
+
+        # Parse fallback branches if provided
+        fallback_branch_ids = (
+            fallback_branches.split(",") if fallback_branches else None
+        )
+
+        # Create service and execute query
+        service = SqlService(profile=profile)
+
+        with SpinnerProgressTracker().track_spinner("Executing SQL query..."):
+            result = service.execute_query(
+                query=query,
+                fallback_branch_ids=fallback_branch_ids,
+                timeout=timeout,
+                format="table" if output_format in ["table", "csv"] else "json",
+            )
+
+        # Save results to file
+        query_results = result.get("results", {})
+        formatter.save_to_file(query_results, output_file, output_format)
+
+        console.print(
+            f"[green]Query executed and results saved to {output_file}[/green]"
+        )
+
+        # Show query metadata
+        if "query_id" in result:
+            console.print(f"[dim]Query ID: {result['query_id']}[/dim]")
+
+    except Exception as e:
+        formatter.print_error(f"Failed to export query results: {e}")
+        raise typer.Exit(1)
+
+
+@app.command("wait")
+def wait_for_query(
+    query_id: str = typer.Argument(..., help="Query ID to wait for"),
+    profile: Optional[str] = typer.Option(
+        None, "--profile", help="Auth profile to use"
+    ),
+    timeout: int = typer.Option(300, "--timeout", help="Maximum wait time in seconds"),
+    output_format: str = typer.Option(
+        "table", "--format", help="Output format for results (table, json, csv)"
+    ),
+    output_file: Optional[Path] = typer.Option(
+        None, "--output", help="Save results to file when completed"
+    ),
+) -> None:
+    """Wait for a query to complete and optionally display results."""
+    console = Console()
+    formatter = OutputFormatter()
+
+    try:
+        service = SqlService(profile=profile)
+
+        with SpinnerProgressTracker().track_spinner("Waiting for query to complete..."):
+            status_result = service.wait_for_completion(query_id, timeout)
+
+        console.print(f"Query ID: [bold]{query_id}[/bold]")
+        console.print(
+            f"Status: [green]{status_result.get('status', 'completed')}[/green]"
+        )
+
+        # Get and display results if requested
+        if output_file or output_format != "table":
+            with SpinnerProgressTracker().track_spinner("Retrieving results..."):
+                result = service.get_query_results(
+                    query_id,
+                    format="table" if output_format in ["table", "csv"] else "json",
+                )
+
+            if output_file:
+                formatter.save_to_file(result, output_file, output_format)
+                console.print(f"[green]Results saved to {output_file}[/green]")
+            else:
+                formatter.display(result, output_format)
+        else:
+            console.print("Use [bold]pltr sql results <query-id>[/bold] to get results")
+
+    except Exception as e:
+        formatter.print_error(f"Failed while waiting for query: {e}")
+        raise typer.Exit(1)

pltr/commands/verify.py

@@ -97,6 +97,7 @@ def verify(
                     "client_id": client_id,
                     "client_secret": client_secret,
                 },
+                timeout=30,
             )
             if token_response.status_code == 200:
                 access_token = token_response.json().get("access_token")
@@ -108,7 +109,7 @@ def verify(
                 raise typer.Exit(1)
 
         # Make the request to /multipass/api/me
-        response = requests.get(url, headers=headers)
+        response = requests.get(url, headers=headers, timeout=30)
 
         if response.status_code == 200:
             user_info = response.json()

pltr/config/aliases.py

@@ -0,0 +1,254 @@
+"""Alias configuration management for pltr-cli."""
+
+import json
+from typing import Dict, List, Optional
+
+from rich import print as rprint
+from rich.table import Table
+
+from pltr.config.settings import Settings
+
+
+class AliasManager:
+    """Manages command aliases for pltr-cli."""
+
+    def __init__(self) -> None:
+        """Initialize the alias manager."""
+        settings = Settings()
+        self.config_dir = settings.config_dir
+        self.aliases_file = self.config_dir / "aliases.json"
+        self.aliases = self._load_aliases()
+
+    def _load_aliases(self) -> Dict[str, str]:
+        """Load aliases from the configuration file.
+
+        Returns:
+            Dictionary mapping alias names to commands
+        """
+        if not self.aliases_file.exists():
+            return {}
+
+        try:
+            with open(self.aliases_file, "r") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, IOError):
+            return {}
+
+    def _save_aliases(self) -> None:
+        """Save aliases to the configuration file."""
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+        with open(self.aliases_file, "w") as f:
+            json.dump(self.aliases, f, indent=2, sort_keys=True)
+
+    def add_alias(self, name: str, command: str) -> bool:
+        """Add a new alias.
+
+        Args:
+            name: Alias name
+            command: Command to alias
+
+        Returns:
+            True if alias was added, False if it already exists
+
+        Raises:
+            ValueError: If the alias would create a circular reference
+        """
+        if name in self.aliases:
+            return False
+
+        # Check for circular references
+        if self._would_create_cycle(name, command):
+            raise ValueError(f"Alias '{name}' would create a circular reference")
+
+        self.aliases[name] = command
+        self._save_aliases()
+        return True
+
+    def remove_alias(self, name: str) -> bool:
+        """Remove an alias.
+
+        Args:
+            name: Alias name to remove
+
+        Returns:
+            True if alias was removed, False if it didn't exist
+        """
+        if name not in self.aliases:
+            return False
+
+        del self.aliases[name]
+        self._save_aliases()
+        return True
+
+    def edit_alias(self, name: str, command: str) -> bool:
+        """Edit an existing alias.
+
+        Args:
+            name: Alias name to edit
+            command: New command for the alias
+
+        Returns:
+            True if alias was edited, False if it doesn't exist
+        """
+        if name not in self.aliases:
+            return False
+
+        # Check for circular references
+        if self._would_create_cycle(name, command):
+            raise ValueError(f"Alias '{name}' would create a circular reference")
+
+        self.aliases[name] = command
+        self._save_aliases()
+        return True
+
+    def get_alias(self, name: str) -> Optional[str]:
+        """Get the command for an alias.
+
+        Args:
+            name: Alias name
+
+        Returns:
+            The aliased command, or None if alias doesn't exist
+        """
+        return self.aliases.get(name)
+
+    def list_aliases(self) -> Dict[str, str]:
+        """Get all aliases.
+
+        Returns:
+            Dictionary of all aliases
+        """
+        return self.aliases.copy()
+
+    def resolve_alias(self, command: str, max_depth: int = 10) -> str:
+        """Resolve an alias to its final command.
+
+        Args:
+            command: Command that might be an alias
+            max_depth: Maximum recursion depth for nested aliases
+
+        Returns:
+            The resolved command
+        """
+        resolved = command
+        depth = 0
+        seen = set()
+
+        while resolved in self.aliases and depth < max_depth:
+            if resolved in seen:
+                # Circular reference detected
+                return command
+
+            seen.add(resolved)
+            resolved = self.aliases[resolved]
+            depth += 1
+
+        return resolved
+
+    def _would_create_cycle(self, name: str, command: str) -> bool:
+        """Check if adding/editing an alias would create a cycle.
+
+        Args:
+            name: Alias name
+            command: Command to check
+
+        Returns:
+            True if this would create a cycle
+        """
+        # Direct self-reference
+        if command == name:
+            return True
+
+        # Follow the chain starting from command
+        # If we ever reach 'name', it would create a cycle
+        current = command
+        visited = set()
+
+        while current in self.aliases:
+            # Check for existing cycles
+            if current in visited:
+                break
+            visited.add(current)
+
+            # Get what this alias points to
+            current = self.aliases[current]
+
+            # If we reached the name we're trying to add, it's a cycle
+            if current == name:
+                return True
+
+        return False
+
+    def display_aliases(self, name: Optional[str] = None) -> None:
+        """Display aliases in a formatted table.
+
+        Args:
+            name: Optional specific alias to display
+        """
+        if name:
+            command = self.get_alias(name)
+            if command:
+                rprint(f"[green]{name}[/green] → {command}")
+            else:
+                rprint(f"[red]Alias '{name}' not found[/red]")
+            return
+
+        if not self.aliases:
+            rprint("[yellow]No aliases configured[/yellow]")
+            return
+
+        table = Table(title="Command Aliases", show_header=True)
+        table.add_column("Alias", style="cyan")
+        table.add_column("Command", style="green")
+
+        for alias_name, command in sorted(self.aliases.items()):
+            table.add_row(alias_name, command)
+
+        rprint(table)
+
+    def get_completion_items(self) -> List[str]:
+        """Get alias names for shell completion.
+
+        Returns:
+            List of alias names
+        """
+        return list(self.aliases.keys())
+
+    def clear_all(self) -> int:
+        """Clear all aliases.
+
+        Returns:
+            Number of aliases cleared
+        """
+        count = len(self.aliases)
+        self.aliases = {}
+        self._save_aliases()
+        return count
+
+    def import_aliases(self, data: Dict[str, str]) -> int:
+        """Import aliases from a dictionary.
+
+        Args:
+            data: Dictionary of aliases to import
+
+        Returns:
+            Number of aliases imported
+        """
+        count = 0
+        for name, command in data.items():
+            if not self._would_create_cycle(name, command):
+                self.aliases[name] = command
+                count += 1
+
+        if count > 0:
+            self._save_aliases()
+
+        return count
+
+    def export_aliases(self) -> Dict[str, str]:
+        """Export all aliases.
+
+        Returns:
+            Dictionary of all aliases
+        """
+        return self.aliases.copy()

pltr/services/__init__.py

@@ -1 +1,5 @@
 """Service wrappers for Foundry SDK."""
+
+from .admin import AdminService
+
+__all__ = ["AdminService"]