pltr-cli 0.1.2__py3-none-any.whl → 0.2.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/services/__init__.py

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

pltr/services/admin.py

@@ -0,0 +1,314 @@
+"""
+Admin service wrapper for Foundry SDK admin operations.
+Provides a high-level interface for user, group, role, and organization management.
+"""
+
+from typing import Any, Dict, Optional
+import json
+
+from .base import BaseService
+
+
+class AdminService(BaseService):
+    """Service wrapper for Foundry admin operations."""
+
+    def _get_service(self) -> Any:
+        """Get the Foundry admin service."""
+        return self.client.admin
+
+    # User Management Methods
+    def list_users(
+        self, page_size: Optional[int] = None, page_token: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        List all users in the organization.
+
+        Args:
+            page_size: Maximum number of users to return per page
+            page_token: Token for pagination (from previous response)
+
+        Returns:
+            Dictionary containing user list and pagination info
+        """
+        try:
+            response = self.service.User.list(
+                page_size=page_size, page_token=page_token
+            )
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to list users: {str(e)}")
+
+    def get_user(self, user_id: str) -> Dict[str, Any]:
+        """
+        Get a specific user by ID.
+
+        Args:
+            user_id: The user ID or RID
+
+        Returns:
+            Dictionary containing user information
+        """
+        try:
+            response = self.service.User.get(user_id)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get user {user_id}: {str(e)}")
+
+    def get_current_user(self) -> Dict[str, Any]:
+        """
+        Get information about the current authenticated user.
+
+        Returns:
+            Dictionary containing current user information
+        """
+        try:
+            response = self.service.User.get_current()
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get current user: {str(e)}")
+
+    def search_users(
+        self,
+        query: str,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Search for users by query string.
+
+        Args:
+            query: Search query string
+            page_size: Maximum number of users to return per page
+            page_token: Token for pagination (from previous response)
+
+        Returns:
+            Dictionary containing search results and pagination info
+        """
+        try:
+            response = self.service.User.search(
+                query=query, page_size=page_size, page_token=page_token
+            )
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to search users: {str(e)}")
+
+    def get_user_markings(self, user_id: str) -> Dict[str, Any]:
+        """
+        Get markings/permissions for a specific user.
+
+        Args:
+            user_id: The user ID or RID
+
+        Returns:
+            Dictionary containing user markings information
+        """
+        try:
+            response = self.service.User.get_markings(user_id)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get user markings for {user_id}: {str(e)}")
+
+    def revoke_user_tokens(self, user_id: str) -> Dict[str, Any]:
+        """
+        Revoke all tokens for a specific user.
+
+        Args:
+            user_id: The user ID or RID
+
+        Returns:
+            Dictionary containing operation result
+        """
+        try:
+            self.service.User.revoke_all_tokens(user_id)
+            return {
+                "success": True,
+                "message": f"All tokens revoked for user {user_id}",
+            }
+        except Exception as e:
+            raise RuntimeError(f"Failed to revoke tokens for user {user_id}: {str(e)}")
+
+    # Group Management Methods
+    def list_groups(
+        self, page_size: Optional[int] = None, page_token: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        List all groups in the organization.
+
+        Args:
+            page_size: Maximum number of groups to return per page
+            page_token: Token for pagination (from previous response)
+
+        Returns:
+            Dictionary containing group list and pagination info
+        """
+        try:
+            response = self.service.Group.list(
+                page_size=page_size, page_token=page_token
+            )
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to list groups: {str(e)}")
+
+    def get_group(self, group_id: str) -> Dict[str, Any]:
+        """
+        Get a specific group by ID.
+
+        Args:
+            group_id: The group ID or RID
+
+        Returns:
+            Dictionary containing group information
+        """
+        try:
+            response = self.service.Group.get(group_id)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get group {group_id}: {str(e)}")
+
+    def search_groups(
+        self,
+        query: str,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Search for groups by query string.
+
+        Args:
+            query: Search query string
+            page_size: Maximum number of groups to return per page
+            page_token: Token for pagination (from previous response)
+
+        Returns:
+            Dictionary containing search results and pagination info
+        """
+        try:
+            response = self.service.Group.search(
+                query=query, page_size=page_size, page_token=page_token
+            )
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to search groups: {str(e)}")
+
+    def create_group(
+        self,
+        name: str,
+        description: Optional[str] = None,
+        organization_rid: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Create a new group.
+
+        Args:
+            name: The group name
+            description: Optional group description
+            organization_rid: Optional organization RID
+
+        Returns:
+            Dictionary containing created group information
+        """
+        try:
+            # Build create request parameters
+            create_params = {"name": name}
+            if description:
+                create_params["description"] = description
+            if organization_rid:
+                create_params["organization_rid"] = organization_rid
+
+            response = self.service.Group.create(**create_params)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to create group '{name}': {str(e)}")
+
+    def delete_group(self, group_id: str) -> Dict[str, Any]:
+        """
+        Delete a specific group.
+
+        Args:
+            group_id: The group ID or RID
+
+        Returns:
+            Dictionary containing operation result
+        """
+        try:
+            self.service.Group.delete(group_id)
+            return {
+                "success": True,
+                "message": f"Group {group_id} deleted successfully",
+            }
+        except Exception as e:
+            raise RuntimeError(f"Failed to delete group {group_id}: {str(e)}")
+
+    # Organization Management Methods
+    def get_organization(self, organization_id: str) -> Dict[str, Any]:
+        """
+        Get organization information.
+
+        Args:
+            organization_id: The organization ID or RID
+
+        Returns:
+            Dictionary containing organization information
+        """
+        try:
+            response = self.service.Organization.get(organization_id)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get organization {organization_id}: {str(e)}"
+            )
+
+    # Role Management Methods
+    def get_role(self, role_id: str) -> Dict[str, Any]:
+        """
+        Get role information.
+
+        Args:
+            role_id: The role ID or RID
+
+        Returns:
+            Dictionary containing role information
+        """
+        try:
+            response = self.service.Role.get(role_id)
+            return self._serialize_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get role {role_id}: {str(e)}")
+
+    def _serialize_response(self, response: Any) -> Dict[str, Any]:
+        """
+        Convert response object to serializable dictionary.
+
+        Args:
+            response: Response object from SDK
+
+        Returns:
+            Serializable dictionary representation
+        """
+        if response is None:
+            return {}
+
+        # Handle different response types
+        if hasattr(response, "dict"):
+            # Pydantic models
+            return response.dict()
+        elif hasattr(response, "__dict__"):
+            # Regular objects
+            result = {}
+            for key, value in response.__dict__.items():
+                if not key.startswith("_"):
+                    try:
+                        # Try to serialize the value
+                        json.dumps(value)
+                        result[key] = value
+                    except (TypeError, ValueError):
+                        # Convert non-serializable values to string
+                        result[key] = str(value)
+            return result
+        else:
+            # Primitive types or already serializable
+            try:
+                json.dumps(response)
+                return response
+            except (TypeError, ValueError):
+                return {"data": str(response)}