pltr-cli 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

pltr/services/sql.py

@@ -0,0 +1,340 @@
+"""
+SQL service wrapper for Foundry SDK SQL queries.
+Provides a high-level interface for executing SQL queries against Foundry datasets.
+"""
+
+import time
+from typing import Any, Dict, List, Optional, Union
+import json
+
+from foundry_sdk.v2.sql_queries.models import (
+    RunningQueryStatus,
+    SucceededQueryStatus,
+    FailedQueryStatus,
+    CanceledQueryStatus,
+)
+
+from .base import BaseService
+
+
+class SqlService(BaseService):
+    """Service wrapper for Foundry SQL query operations."""
+
+    def _get_service(self) -> Any:
+        """Get the Foundry SQL queries service."""
+        return self.client.sql_queries.SqlQuery
+
+    def execute_query(
+        self,
+        query: str,
+        fallback_branch_ids: Optional[List[str]] = None,
+        timeout: int = 300,
+        format: str = "table",
+    ) -> Dict[str, Any]:
+        """
+        Execute a SQL query and wait for completion.
+
+        Args:
+            query: SQL query string
+            fallback_branch_ids: Optional list of branch IDs for fallback
+            timeout: Maximum time to wait for query completion (seconds)
+            format: Output format for results ('table', 'json', 'raw')
+
+        Returns:
+            Dictionary containing query results and metadata
+
+        Raises:
+            RuntimeError: If query execution fails or times out
+        """
+        try:
+            # Submit the query
+            status = self.service.execute(
+                query=query, fallback_branch_ids=fallback_branch_ids
+            )
+
+            # If the query completed immediately
+            if isinstance(status, SucceededQueryStatus):
+                return self._format_completed_query(status.query_id, format)
+            elif isinstance(status, FailedQueryStatus):
+                raise RuntimeError(f"Query failed: {status.error_message}")
+            elif isinstance(status, CanceledQueryStatus):
+                raise RuntimeError("Query was canceled")
+            elif isinstance(status, RunningQueryStatus):
+                # Wait for completion
+                return self._wait_for_query_completion(status.query_id, timeout, format)
+            else:
+                raise RuntimeError(f"Unknown query status type: {type(status)}")
+
+        except Exception as e:
+            if isinstance(e, RuntimeError):
+                raise
+            raise RuntimeError(f"Failed to execute query: {e}")
+
+    def submit_query(
+        self, query: str, fallback_branch_ids: Optional[List[str]] = None
+    ) -> Dict[str, Any]:
+        """
+        Submit a SQL query without waiting for completion.
+
+        Args:
+            query: SQL query string
+            fallback_branch_ids: Optional list of branch IDs for fallback
+
+        Returns:
+            Dictionary containing query ID and initial status
+
+        Raises:
+            RuntimeError: If query submission fails
+        """
+        try:
+            status = self.service.execute(
+                query=query, fallback_branch_ids=fallback_branch_ids
+            )
+            return self._format_query_status(status)
+        except Exception as e:
+            raise RuntimeError(f"Failed to submit query: {e}")
+
+    def get_query_status(self, query_id: str) -> Dict[str, Any]:
+        """
+        Get the status of a submitted query.
+
+        Args:
+            query_id: Query identifier
+
+        Returns:
+            Dictionary containing query status information
+
+        Raises:
+            RuntimeError: If status check fails
+        """
+        try:
+            status = self.service.get_status(query_id)
+            return self._format_query_status(status)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get query status: {e}")
+
+    def get_query_results(self, query_id: str, format: str = "table") -> Dict[str, Any]:
+        """
+        Get the results of a completed query.
+
+        Args:
+            query_id: Query identifier
+            format: Output format ('table', 'json', 'raw')
+
+        Returns:
+            Dictionary containing query results
+
+        Raises:
+            RuntimeError: If results retrieval fails
+        """
+        try:
+            # First check if the query has completed successfully
+            status = self.service.get_status(query_id)
+            if not isinstance(status, SucceededQueryStatus):
+                status_info = self._format_query_status(status)
+                if isinstance(status, FailedQueryStatus):
+                    raise RuntimeError(f"Query failed: {status.error_message}")
+                elif isinstance(status, CanceledQueryStatus):
+                    raise RuntimeError("Query was canceled")
+                elif isinstance(status, RunningQueryStatus):
+                    raise RuntimeError("Query is still running")
+                else:
+                    raise RuntimeError(f"Query status: {status_info['status']}")
+
+            # Get the results
+            results_bytes = self.service.get_results(query_id)
+            return self._format_query_results(results_bytes, format)
+
+        except Exception as e:
+            if isinstance(e, RuntimeError):
+                raise
+            raise RuntimeError(f"Failed to get query results: {e}")
+
+    def cancel_query(self, query_id: str) -> Dict[str, Any]:
+        """
+        Cancel a running query.
+
+        Args:
+            query_id: Query identifier
+
+        Returns:
+            Dictionary containing cancellation status
+
+        Raises:
+            RuntimeError: If cancellation fails
+        """
+        try:
+            self.service.cancel(query_id)
+            # Get updated status after cancellation
+            status = self.service.get_status(query_id)
+            return self._format_query_status(status)
+        except Exception as e:
+            raise RuntimeError(f"Failed to cancel query: {e}")
+
+    def wait_for_completion(
+        self, query_id: str, timeout: int = 300, poll_interval: int = 2
+    ) -> Dict[str, Any]:
+        """
+        Wait for a query to complete.
+
+        Args:
+            query_id: Query identifier
+            timeout: Maximum time to wait (seconds)
+            poll_interval: Time between status checks (seconds)
+
+        Returns:
+            Dictionary containing final query status
+
+        Raises:
+            RuntimeError: If query fails or times out
+        """
+        start_time = time.time()
+
+        while time.time() - start_time < timeout:
+            try:
+                status = self.service.get_status(query_id)
+
+                if isinstance(status, SucceededQueryStatus):
+                    return self._format_query_status(status)
+                elif isinstance(status, FailedQueryStatus):
+                    raise RuntimeError(f"Query failed: {status.error_message}")
+                elif isinstance(status, CanceledQueryStatus):
+                    raise RuntimeError("Query was canceled")
+                elif isinstance(status, RunningQueryStatus):
+                    # Still running, continue waiting
+                    time.sleep(poll_interval)
+                    continue
+                else:
+                    raise RuntimeError(f"Unknown status type: {type(status)}")
+
+            except Exception as e:
+                if isinstance(e, RuntimeError):
+                    raise
+                raise RuntimeError(f"Error checking query status: {e}")
+
+        # Timeout reached
+        raise RuntimeError(f"Query timed out after {timeout} seconds")
+
+    def _wait_for_query_completion(
+        self, query_id: str, timeout: int, format: str
+    ) -> Dict[str, Any]:
+        """
+        Wait for query completion and return formatted results.
+
+        Args:
+            query_id: Query identifier
+            timeout: Maximum wait time
+            format: Result format
+
+        Returns:
+            Dictionary with query results
+        """
+        # Wait for completion
+        self.wait_for_completion(query_id, timeout)
+
+        # Get results
+        return self._format_completed_query(query_id, format)
+
+    def _format_completed_query(self, query_id: str, format: str) -> Dict[str, Any]:
+        """
+        Format a completed query's results.
+
+        Args:
+            query_id: Query identifier
+            format: Result format
+
+        Returns:
+            Formatted query results
+        """
+        results_bytes = self.service.get_results(query_id)
+        results = self._format_query_results(results_bytes, format)
+
+        return {
+            "query_id": query_id,
+            "status": "succeeded",
+            "results": results,
+        }
+
+    def _format_query_status(
+        self,
+        status: Union[
+            RunningQueryStatus,
+            SucceededQueryStatus,
+            FailedQueryStatus,
+            CanceledQueryStatus,
+        ],
+    ) -> Dict[str, Any]:
+        """
+        Format query status for consistent output.
+
+        Args:
+            status: Query status object
+
+        Returns:
+            Formatted status dictionary
+        """
+        base_info: Dict[str, Any] = {"status": status.type}
+
+        if isinstance(status, (RunningQueryStatus, SucceededQueryStatus)):
+            base_info["query_id"] = status.query_id
+        elif isinstance(status, FailedQueryStatus):
+            base_info["error_message"] = status.error_message
+
+        return base_info
+
+    def _format_query_results(self, results_bytes: bytes, format: str) -> Any:
+        """
+        Format query results based on the requested format.
+
+        Args:
+            results_bytes: Raw results from the API
+            format: Desired output format
+
+        Returns:
+            Formatted results
+        """
+        if format == "raw":
+            return results_bytes
+
+        # Try to decode as text first
+        try:
+            results_text = results_bytes.decode("utf-8")
+        except UnicodeDecodeError:
+            # If it's binary data, return as base64 or hex
+            return {
+                "type": "binary",
+                "size_bytes": len(results_bytes),
+                "data": results_bytes.hex()[:200] + "..."
+                if len(results_bytes) > 100
+                else results_bytes.hex(),
+            }
+
+        if format == "json":
+            try:
+                # Try to parse as JSON
+                return json.loads(results_text)
+            except json.JSONDecodeError:
+                # Return as text if not valid JSON
+                return {"text": results_text}
+
+        elif format == "table":
+            # For table format, we'll return structured data
+            # that the formatter can convert to a table
+            try:
+                # Try parsing as JSON first for structured data
+                data = json.loads(results_text)
+                if isinstance(data, list) and data and isinstance(data[0], dict):
+                    # List of dictionaries - perfect for table format
+                    return data
+                else:
+                    return {"result": data}
+            except json.JSONDecodeError:
+                # Return as text data
+                lines = results_text.strip().split("\n")
+                if len(lines) == 1:
+                    return {"result": lines[0]}
+                else:
+                    return {"results": lines}
+
+        # Default: return as text
+        return {"text": results_text}

pltr/utils/completion.py

@@ -0,0 +1,170 @@
+"""Shell completion utilities for pltr CLI."""
+
+import os
+from typing import List
+from pathlib import Path
+import json
+
+from pltr.config.profiles import ProfileManager
+
+
+def get_cached_rids() -> List[str]:
+    """Get recently used RIDs from cache."""
+    cache_dir = Path.home() / ".cache" / "pltr"
+    rid_cache_file = cache_dir / "recent_rids.json"
+
+    if rid_cache_file.exists():
+        try:
+            with open(rid_cache_file) as f:
+                data = json.load(f)
+                return data.get("rids", [])
+        except Exception:
+            pass
+
+    # Return some example RIDs if no cache
+    return [
+        "ri.foundry.main.dataset.",
+        "ri.foundry.main.folder.",
+        "ri.foundry.main.ontology.",
+    ]
+
+
+def cache_rid(rid: str):
+    """Cache a RID for future completions."""
+    cache_dir = Path.home() / ".cache" / "pltr"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    rid_cache_file = cache_dir / "recent_rids.json"
+
+    # Load existing cache
+    rids = []
+    if rid_cache_file.exists():
+        try:
+            with open(rid_cache_file) as f:
+                data = json.load(f)
+                rids = data.get("rids", [])
+        except Exception:
+            pass
+
+    # Add new RID (keep last 50)
+    if rid not in rids:
+        rids.insert(0, rid)
+        rids = rids[:50]
+
+    # Save cache
+    try:
+        with open(rid_cache_file, "w") as f:
+            json.dump({"rids": rids}, f)
+    except Exception:
+        pass
+
+
+def complete_rid(incomplete: str):
+    """Complete RID arguments."""
+    rids = get_cached_rids()
+    return [rid for rid in rids if rid.startswith(incomplete)]
+
+
+def complete_profile(incomplete: str):
+    """Complete profile names."""
+    try:
+        manager = ProfileManager()
+        profiles = manager.list_profiles()
+        return [profile for profile in profiles if profile.startswith(incomplete)]
+    except Exception:
+        return []
+
+
+def complete_output_format(incomplete: str):
+    """Complete output format options."""
+    formats = ["table", "json", "csv"]
+    return [fmt for fmt in formats if fmt.startswith(incomplete)]
+
+
+def complete_sql_query(incomplete: str):
+    """Complete SQL query templates."""
+    templates = [
+        "SELECT * FROM ",
+        "SELECT COUNT(*) FROM ",
+        "SELECT DISTINCT ",
+        "WHERE ",
+        "GROUP BY ",
+        "ORDER BY ",
+        "LIMIT 10",
+        "JOIN ",
+        "LEFT JOIN ",
+        "INNER JOIN ",
+    ]
+    return [tmpl for tmpl in templates if tmpl.lower().startswith(incomplete.lower())]
+
+
+def complete_ontology_action(incomplete: str):
+    """Complete ontology action names."""
+    # This would ideally fetch from the API but for now return common patterns
+    actions = [
+        "create",
+        "update",
+        "delete",
+        "createOrUpdate",
+        "link",
+        "unlink",
+    ]
+    return [action for action in actions if action.startswith(incomplete)]
+
+
+def complete_file_path(incomplete: str):
+    """Complete file paths."""
+    # This is handled by shell natively, but we can provide hints
+    path = Path(incomplete) if incomplete else Path.cwd()
+
+    if incomplete and not path.exists():
+        parent = path.parent
+        prefix = path.name
+    else:
+        parent = path if path.is_dir() else path.parent
+        prefix = ""
+
+    try:
+        items = []
+        for item in parent.iterdir():
+            if item.name.startswith(prefix):
+                # Return path strings - shell will handle directory indicators
+                items.append(str(item))
+        return items
+    except Exception:
+        return []
+
+
+def setup_completion_environment():
+    """Set up environment for shell completion support."""
+    # This is called when the CLI starts to register completion handlers
+
+    # Check if we're in completion mode
+    if os.environ.get("_PLTR_COMPLETE"):
+        # We're generating completions
+        # Set up any necessary context
+        pass
+
+
+def handle_completion():
+    """Handle shell completion requests."""
+    # This is the main entry point for completion handling
+    # It's called when _PLTR_COMPLETE environment variable is set
+
+    complete_var = os.environ.get("_PLTR_COMPLETE")
+    if not complete_var:
+        return False
+
+    # Click handles the completion automatically through Typer
+    # Our custom completion functions are registered via autocompletion parameter
+    return True
+
+
+# Register custom completion functions for specific parameter types
+COMPLETION_FUNCTIONS = {
+    "rid": complete_rid,
+    "profile": complete_profile,
+    "output_format": complete_output_format,
+    "sql_query": complete_sql_query,
+    "ontology_action": complete_ontology_action,
+    "file_path": complete_file_path,
+}