pltr-cli 0.10.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

pltr/services/dataset.py

@@ -2,10 +2,12 @@
 Dataset service wrapper for Foundry SDK.
 """
 
-from typing import Any, Optional, List, Dict, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 from pathlib import Path
 import csv
 
+from ..config.settings import Settings
+from ..utils.pagination import PaginationConfig, PaginationResult
 from .base import BaseService
 
 
@@ -312,6 +314,29 @@ class DatasetService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to read dataset {dataset_rid}: {e}")
 
+    def preview_data(
+        self,
+        dataset_rid: str,
+        limit: int = 10,
+    ) -> List[Dict[str, Any]]:
+        """
+        Preview dataset contents as a list of records.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            limit: Maximum number of rows to return
+
+        Returns:
+            List of dictionaries representing rows
+        """
+        try:
+            # Use read_table with pandas format for easy conversion
+            df = self.read_table(dataset_rid, format="pandas")
+            # Limit rows and convert to records
+            return df.head(limit).to_dict(orient="records")
+        except Exception as e:
+            raise RuntimeError(f"Failed to preview dataset {dataset_rid}: {e}")
+
     def delete_dataset(self, dataset_rid: str) -> bool:
         """
         Delete a dataset.
@@ -487,18 +512,15 @@ class DatasetService(BaseService):
             # Ensure output directory exists
             output_path.parent.mkdir(parents=True, exist_ok=True)
 
-            file_content = self.service.Dataset.File.read(
+            # Use Dataset.File.content() which returns bytes directly
+            # Note: In SDK v1.27.0, the method is 'content' not 'read'
+            file_content = self.service.Dataset.File.content(
                 dataset_rid=dataset_rid, file_path=file_path, branch_name=branch
             )
 
-            # Write file content to disk
+            # Write file content to disk (file_content is bytes)
             with open(output_path, "wb") as f:
-                if hasattr(file_content, "read"):
-                    # If it's a stream
-                    f.write(file_content.read())
-                else:
-                    # If it's bytes
-                    f.write(file_content)
+                f.write(file_content)
 
             return {
                 "dataset_rid": dataset_rid,
@@ -519,6 +541,8 @@ class DatasetService(BaseService):
         """
         List files in a dataset.
 
+        DEPRECATED: Use list_files_paginated() instead for better pagination support.
+
         Args:
             dataset_rid: Dataset Resource Identifier
             branch: Dataset branch name
@@ -543,6 +567,53 @@ class DatasetService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to list files in dataset {dataset_rid}: {e}")
 
+    def list_files_paginated(
+        self,
+        dataset_rid: str,
+        branch: str,
+        config: PaginationConfig,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> PaginationResult:
+        """
+        List files with full pagination control.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            branch: Dataset branch name
+            config: Pagination configuration
+            progress_callback: Optional progress callback
+
+        Returns:
+            PaginationResult with file information and metadata
+        """
+        try:
+            settings = Settings()
+
+            # Get iterator from SDK - ResourceIterator with next_page_token support
+            iterator = self.service.Dataset.File.list(
+                dataset_rid=dataset_rid,
+                branch_name=branch,
+                page_size=config.page_size or settings.get("page_size", 20),
+            )
+
+            # Use iterator pagination handler
+            result = self._paginate_iterator(iterator, config, progress_callback)
+
+            # Format file information
+            result.data = [
+                {
+                    "path": file.path,
+                    "size_bytes": getattr(file, "size_bytes", None),
+                    "last_modified": getattr(file, "last_modified", None),
+                    "transaction_rid": getattr(file, "transaction_rid", None),
+                }
+                for file in result.data
+            ]
+
+            return result
+        except Exception as e:
+            raise RuntimeError(f"Failed to list files: {e}")
+
     def get_branches(self, dataset_rid: str) -> List[Dict[str, Any]]:
         """
         Get list of branches for a dataset.

pltr/services/mediasets.py

@@ -226,15 +226,7 @@ class MediaSetsService(BaseService):
                     preview=preview,
                 )
 
-            with open(output_path_obj, "wb") as file:
-                if hasattr(response, "content"):
-                    file.write(response.content)
-                else:
-                    # Handle streaming response
-                    for chunk in response:
-                        file.write(chunk)
-
-            file_size = output_path_obj.stat().st_size
+            file_size = self._write_response_to_file(response, output_path_obj)
             return {
                 "media_set_rid": media_set_rid,
                 "media_item_rid": media_item_rid,
@@ -291,3 +283,146 @@ class MediaSetsService(BaseService):
             "url": getattr(reference_response, "url", "unknown"),
             "expires_at": getattr(reference_response, "expires_at", None),
         }
+
+    def _write_response_to_file(self, response: Any, output_path: Path) -> int:
+        """
+        Write response content to file and return file size.
+
+        Args:
+            response: SDK response object (with .content attribute or iterable)
+            output_path: Path object for the output file
+
+        Returns:
+            File size in bytes
+        """
+        with open(output_path, "wb") as file:
+            if hasattr(response, "content"):
+                file.write(response.content)
+            else:
+                # Handle streaming response
+                for chunk in response:
+                    file.write(chunk)
+        return output_path.stat().st_size
+
+    def calculate_thumbnail(
+        self,
+        media_set_rid: str,
+        media_item_rid: str,
+        preview: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Initiate thumbnail generation for an image.
+
+        Args:
+            media_set_rid: Media Set Resource Identifier
+            media_item_rid: Media Item Resource Identifier
+            preview: Enable preview mode
+
+        Returns:
+            Thumbnail calculation status information
+        """
+        try:
+            response = self.service.MediaSet.calculate(
+                media_set_rid=media_set_rid,
+                media_item_rid=media_item_rid,
+                preview=preview,
+            )
+            return self._format_thumbnail_status(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to calculate thumbnail: {e}")
+
+    def retrieve_thumbnail(
+        self,
+        media_set_rid: str,
+        media_item_rid: str,
+        output_path: str,
+        preview: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Retrieve a calculated thumbnail (200px wide webp).
+
+        Args:
+            media_set_rid: Media Set Resource Identifier
+            media_item_rid: Media Item Resource Identifier
+            output_path: Local path where thumbnail should be saved
+            preview: Enable preview mode
+
+        Returns:
+            Download response information
+        """
+        try:
+            output_path_obj = Path(output_path)
+            output_path_obj.parent.mkdir(parents=True, exist_ok=True)
+
+            response = self.service.MediaSet.retrieve(
+                media_set_rid=media_set_rid,
+                media_item_rid=media_item_rid,
+                preview=preview,
+            )
+
+            file_size = self._write_response_to_file(response, output_path_obj)
+
+            # Validate that we received actual content
+            if file_size == 0:
+                output_path_obj.unlink(missing_ok=True)
+                raise RuntimeError(
+                    "Downloaded thumbnail is empty - thumbnail may not be ready yet"
+                )
+
+            return {
+                "media_set_rid": media_set_rid,
+                "media_item_rid": media_item_rid,
+                "output_path": str(output_path_obj),
+                "file_size": file_size,
+                "downloaded": True,
+                "format": "image/webp",
+            }
+        except Exception as e:
+            raise RuntimeError(f"Failed to retrieve thumbnail: {e}")
+
+    def upload_temp_media(
+        self,
+        file_path: str,
+        filename: Optional[str] = None,
+        attribution: Optional[str] = None,
+        preview: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Upload temporary media (auto-deleted after 1 hour if not persisted).
+
+        Args:
+            file_path: Local path to the file to upload
+            filename: Optional filename override
+            attribution: Optional attribution string
+            preview: Enable preview mode
+
+        Returns:
+            Media reference information
+        """
+        try:
+            file_path_obj = Path(file_path)
+            if not file_path_obj.exists():
+                raise FileNotFoundError(f"File not found: {file_path}")
+
+            # Use provided filename or default to file name
+            upload_filename = filename or file_path_obj.name
+
+            with open(file_path_obj, "rb") as file:
+                response = self.service.MediaSet.upload_media(
+                    body=file,
+                    filename=upload_filename,
+                    attribution=attribution,
+                    preview=preview,
+                )
+
+            return self._format_media_reference(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to upload temporary media: {e}")
+
+    def _format_thumbnail_status(self, status_response: Any) -> Dict[str, Any]:
+        """Format thumbnail calculation status response for display."""
+        return {
+            "status": getattr(status_response, "status", "unknown"),
+            "transformation_id": getattr(status_response, "transformation_id", None),
+            "media_item_rid": getattr(status_response, "media_item_rid", None),
+        }

pltr/services/ontology.py

@@ -2,7 +2,10 @@
 Ontology service wrappers for Foundry SDK.
 """
 
-from typing import Any, Optional, Dict, List
+from typing import Any, Callable, Dict, List, Optional
+
+from ..config.settings import Settings
+from ..utils.pagination import PaginationConfig, PaginationResult
 from .base import BaseService
 
 
@@ -165,6 +168,8 @@ class OntologyObjectService(BaseService):
         """
         List objects of a specific type.
 
+        DEPRECATED: Use list_objects_paginated() instead for better pagination support.
+
         Args:
             ontology_rid: Ontology Resource Identifier
             object_type: Object type API name
@@ -188,6 +193,48 @@ class OntologyObjectService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to list objects: {e}")
 
+    def list_objects_paginated(
+        self,
+        ontology_rid: str,
+        object_type: str,
+        config: PaginationConfig,
+        properties: Optional[List[str]] = None,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+    ) -> PaginationResult:
+        """
+        List objects with full pagination control.
+
+        Args:
+            ontology_rid: Ontology Resource Identifier
+            object_type: Object type API name
+            config: Pagination configuration
+            properties: List of properties to include
+            progress_callback: Optional progress callback
+
+        Returns:
+            PaginationResult with objects and metadata
+        """
+        try:
+            settings = Settings()
+
+            # Get iterator from SDK - ResourceIterator with next_page_token support
+            iterator = self.service.OntologyObject.list(
+                ontology_rid,
+                object_type,
+                page_size=config.page_size or settings.get("page_size", 20),
+                properties=properties,
+            )
+
+            # Use iterator pagination handler
+            result = self._paginate_iterator(iterator, config, progress_callback)
+
+            # Format objects
+            result.data = [self._format_object(obj) for obj in result.data]
+
+            return result
+        except Exception as e:
+            raise RuntimeError(f"Failed to list objects: {e}")
+
     def get_object(
         self,
         ontology_rid: str,
@@ -287,6 +334,77 @@ class OntologyObjectService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to list linked objects: {e}")
 
+    def count_objects(
+        self,
+        ontology_rid: str,
+        object_type: str,
+        branch: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Count objects of a specific type.
+
+        Args:
+            ontology_rid: Ontology Resource Identifier
+            object_type: Object type API name
+            branch: Branch name (optional)
+
+        Returns:
+            Dictionary containing count information
+        """
+        try:
+            count = self.service.OntologyObject.count(
+                ontology_rid,
+                object_type,
+                branch_name=branch,
+            )
+            return {
+                "ontology_rid": ontology_rid,
+                "object_type": object_type,
+                "count": count,
+                "branch": branch,
+            }
+        except Exception as e:
+            raise RuntimeError(f"Failed to count objects: {e}")
+
+    def search_objects(
+        self,
+        ontology_rid: str,
+        object_type: str,
+        query: str,
+        page_size: Optional[int] = None,
+        properties: Optional[List[str]] = None,
+        branch: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Search objects by query.
+
+        Args:
+            ontology_rid: Ontology Resource Identifier
+            object_type: Object type API name
+            query: Search query string
+            page_size: Number of results per page
+            properties: List of properties to include
+            branch: Branch name (optional)
+
+        Returns:
+            List of matching object dictionaries
+        """
+        try:
+            result = self.service.OntologyObject.search(
+                ontology_rid,
+                object_type,
+                query=query,
+                page_size=page_size,
+                properties=properties,
+                branch_name=branch,
+            )
+            objects = []
+            for obj in result:
+                objects.append(self._format_object(obj))
+            return objects
+        except Exception as e:
+            raise RuntimeError(f"Failed to search objects: {e}")
+
     def _format_object(self, obj: Any) -> Dict[str, Any]:
         """Format object for consistent output."""
         # Objects may have various properties - extract them dynamically

pltr/services/orchestration.py

@@ -3,8 +3,10 @@ Orchestration service wrapper for Foundry SDK v2 API.
 Provides operations for managing builds, jobs, and schedules.
 """
 
-from typing import Any, Optional, Dict, List
+from typing import Any, Callable, Dict, List, Optional
 
+from ..config.settings import Settings
+from ..utils.pagination import PaginationConfig, PaginationResult
 from .base import BaseService
 
 
@@ -134,6 +136,8 @@ class OrchestrationService(BaseService):
         """
         Search for builds.
 
+        DEPRECATED: Use search_builds_paginated() instead for better pagination support.
+
         Args:
             page_size: Number of results per page
             page_token: Token for pagination
@@ -155,6 +159,63 @@ class OrchestrationService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to search builds: {e}")
 
+    def search_builds_paginated(
+        self,
+        config: PaginationConfig,
+        progress_callback: Optional[Callable[[int, int], None]] = None,
+        **search_params,
+    ) -> PaginationResult:
+        """
+        Search for builds with full pagination control.
+
+        Args:
+            config: Pagination configuration
+            progress_callback: Optional progress callback
+            **search_params: Additional search parameters
+
+        Returns:
+            PaginationResult with builds and metadata
+        """
+        try:
+            settings = Settings()
+
+            def fetch_page(page_token: Optional[str]) -> Dict[str, Any]:
+                """Fetch a single page of builds."""
+                kwargs: Dict[str, Any] = {
+                    "page_size": config.page_size or settings.get("page_size", 20),
+                }
+                if page_token:
+                    kwargs["page_token"] = page_token
+                kwargs.update(search_params)
+
+                response = self.service.Build.search(**kwargs)
+                return self._format_builds_search_response(response)
+
+            return self._paginate_response(fetch_page, config, progress_callback)
+        except Exception as e:
+            raise RuntimeError(f"Failed to search builds: {e}")
+
+    def get_builds_batch(self, build_rids: List[str]) -> Dict[str, Any]:
+        """
+        Get multiple builds in batch.
+
+        Args:
+            build_rids: List of Build Resource Identifiers (max 100)
+
+        Returns:
+            Batch response with build information
+        """
+        if len(build_rids) > 100:
+            raise ValueError("Maximum batch size is 100 builds")
+
+        try:
+            # SDK expects list of {"rid": ...} objects for batch operations
+            body = [{"rid": rid} for rid in build_rids]
+            response = self.service.Build.get_batch(body)
+            return self._format_builds_batch_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get builds batch: {e}")
+
     # Job operations
     def get_job(self, job_rid: str) -> Dict[str, Any]:
         """
@@ -353,6 +414,35 @@ class OrchestrationService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to replace schedule {schedule_rid}: {e}")
 
+    def get_schedule_runs(
+        self,
+        schedule_rid: str,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Get recent execution runs for a schedule.
+
+        Args:
+            schedule_rid: Schedule Resource Identifier
+            page_size: Number of results per page
+            page_token: Token for pagination
+
+        Returns:
+            Runs list with pagination info
+        """
+        try:
+            kwargs: Dict[str, Any] = {"schedule_rid": schedule_rid}
+            if page_size is not None:
+                kwargs["page_size"] = page_size
+            if page_token is not None:
+                kwargs["page_token"] = page_token
+
+            response = self.service.Schedule.runs(**kwargs)
+            return self._format_schedule_runs_response(response)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get runs for schedule {schedule_rid}: {e}")
+
     # Formatting methods
     def _format_build_info(self, build: Any) -> Dict[str, Any]:
         """Format build information for consistent output."""
@@ -455,3 +545,45 @@ class OrchestrationService(BaseService):
                     result["jobs"].append(self._format_job_info(item.data))
 
         return result
+
+    def _format_builds_batch_response(self, response: Any) -> Dict[str, Any]:
+        """Format builds batch response."""
+        result: Dict[str, Any] = {"builds": []}
+
+        if hasattr(response, "data"):
+            for item in response.data:
+                if hasattr(item, "data"):
+                    result["builds"].append(self._format_build_info(item.data))
+
+        return result
+
+    def _format_run_info(self, run: Any) -> Dict[str, Any]:
+        """Format schedule run information for consistent output."""
+        info = {}
+
+        for attr in [
+            "rid",
+            "schedule_rid",
+            "status",
+            "created_time",
+            "started_time",
+            "finished_time",
+            "build_rid",
+            "result",
+        ]:
+            if hasattr(run, attr):
+                info[attr] = getattr(run, attr)
+
+        return info
+
+    def _format_schedule_runs_response(self, response: Any) -> Dict[str, Any]:
+        """Format schedule runs response."""
+        result: Dict[str, Any] = {"runs": []}
+
+        if hasattr(response, "data"):
+            result["runs"] = [self._format_run_info(run) for run in response.data]
+
+        if hasattr(response, "next_page_token"):
+            result["next_page_token"] = response.next_page_token
+
+        return result