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

pltr/services/dataset.py

@@ -1,20 +1,22 @@
 """
-Dataset service wrapper for Foundry SDK v2 API.
-Only includes operations that are actually supported by foundry-platform-sdk v1.27.0.
+Dataset service wrapper for Foundry SDK.
 """
 
-from typing import Any, Optional, Dict
+from typing import Any, Optional, List, Dict, Union
+from pathlib import Path
 
 from .base import BaseService
 
 
 class DatasetService(BaseService):
-    """Service wrapper for Foundry dataset operations using v2 API."""
+    """Service wrapper for Foundry dataset operations."""
 
     def _get_service(self) -> Any:
         """Get the Foundry datasets service."""
         return self.client.datasets
 
+    # list_datasets method removed - not supported by foundry-platform-sdk v1.27.0
+
     def get_dataset(self, dataset_rid: str) -> Dict[str, Any]:
         """
         Get information about a specific dataset.
@@ -32,7 +34,26 @@ class DatasetService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to get dataset {dataset_rid}: {e}")
 
-    # get_schema method removed - uses preview-only API
+    def get_schema(self, dataset_rid: str) -> Dict[str, Any]:
+        """
+        Get dataset schema.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+
+        Returns:
+            Schema information
+        """
+        try:
+            schema = self.service.Dataset.get_schema(dataset_rid)
+            return {
+                "dataset_rid": dataset_rid,
+                "schema": schema,
+                "type": str(type(schema)),
+                "status": "Schema retrieved successfully",
+            }
+        except Exception as e:
+            raise RuntimeError(f"Failed to get schema for dataset {dataset_rid}: {e}")
 
     def create_dataset(
         self, name: str, parent_folder_rid: Optional[str] = None
@@ -48,11 +69,20 @@ class DatasetService(BaseService):
             Created dataset information
         """
         try:
-            # The create method parameters depend on the SDK version
+            # Try the v2 API first (Dataset.create)
             dataset = self.service.Dataset.create(
                 name=name, parent_folder_rid=parent_folder_rid
             )
             return self._format_dataset_info(dataset)
+        except AttributeError:
+            # Fallback to service.create_dataset
+            try:
+                dataset = self.service.create_dataset(
+                    name=name, parent_folder_rid=parent_folder_rid
+                )
+                return self._format_dataset_info(dataset)
+            except Exception as e:
+                raise RuntimeError(f"Failed to create dataset '{name}': {e}")
         except Exception as e:
             raise RuntimeError(f"Failed to create dataset '{name}': {e}")
 
@@ -68,10 +98,328 @@ class DatasetService(BaseService):
             Table data in specified format
         """
         try:
+            # Try the v2 API first (Dataset.read_table)
             return self.service.Dataset.read_table(dataset_rid, format=format)
+        except AttributeError:
+            # Fallback to service.read_table
+            try:
+                return self.service.read_table(dataset_rid, format=format)
+            except Exception as e:
+                raise RuntimeError(f"Failed to read dataset {dataset_rid}: {e}")
         except Exception as e:
             raise RuntimeError(f"Failed to read dataset {dataset_rid}: {e}")
 
+    def delete_dataset(self, dataset_rid: str) -> bool:
+        """
+        Delete a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+
+        Returns:
+            True if deletion was successful
+        """
+        try:
+            self.service.delete_dataset(dataset_rid)
+            return True
+        except Exception as e:
+            raise RuntimeError(f"Failed to delete dataset {dataset_rid}: {e}")
+
+    def upload_file(
+        self,
+        dataset_rid: str,
+        file_path: Union[str, Path],
+        branch: str = "master",
+        transaction_rid: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Upload a file to a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            file_path: Path to file to upload
+            branch: Dataset branch name
+            transaction_rid: Transaction RID (optional)
+
+        Returns:
+            Upload result information
+        """
+        file_path = Path(file_path)
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        try:
+            with open(file_path, "rb") as f:
+                result = self.service.upload_file(
+                    dataset_rid=dataset_rid,
+                    file_path=file_path.name,
+                    file_data=f,
+                    branch=branch,
+                    transaction_rid=transaction_rid,
+                )
+
+            return {
+                "dataset_rid": dataset_rid,
+                "file_path": str(file_path),
+                "branch": branch,
+                "size_bytes": file_path.stat().st_size,
+                "uploaded": True,
+                "transaction_rid": getattr(result, "transaction_rid", transaction_rid),
+            }
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to upload file {file_path} to dataset {dataset_rid}: {e}"
+            )
+
+    def download_file(
+        self,
+        dataset_rid: str,
+        file_path: str,
+        output_path: Union[str, Path],
+        branch: str = "master",
+    ) -> Dict[str, Any]:
+        """
+        Download a file from a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            file_path: Path of file within dataset
+            output_path: Local path to save the downloaded file
+            branch: Dataset branch name
+
+        Returns:
+            Download result information
+        """
+        output_path = Path(output_path)
+
+        try:
+            # Ensure output directory exists
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+
+            file_content = self.service.download_file(
+                dataset_rid=dataset_rid, file_path=file_path, branch=branch
+            )
+
+            # Write file content to disk
+            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)
+
+            return {
+                "dataset_rid": dataset_rid,
+                "file_path": file_path,
+                "output_path": str(output_path),
+                "branch": branch,
+                "size_bytes": output_path.stat().st_size,
+                "downloaded": True,
+            }
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to download file {file_path} from dataset {dataset_rid}: {e}"
+            )
+
+    def list_files(
+        self, dataset_rid: str, branch: str = "master"
+    ) -> List[Dict[str, Any]]:
+        """
+        List files in a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            branch: Dataset branch name
+
+        Returns:
+            List of file information dictionaries
+        """
+        try:
+            files = self.service.list_files(dataset_rid=dataset_rid, branch=branch)
+
+            return [
+                {
+                    "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 files
+            ]
+        except Exception as e:
+            raise RuntimeError(f"Failed to list files in dataset {dataset_rid}: {e}")
+
+    def get_branches(self, dataset_rid: str) -> List[Dict[str, Any]]:
+        """
+        Get list of branches for a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+
+        Returns:
+            List of branch information dictionaries
+        """
+        try:
+            branches = self.service.list_branches(dataset_rid=dataset_rid)
+
+            return [
+                {
+                    "name": branch.name,
+                    "transaction_rid": getattr(branch, "transaction_rid", None),
+                    "created_time": getattr(branch, "created_time", None),
+                    "created_by": getattr(branch, "created_by", None),
+                }
+                for branch in branches
+            ]
+        except Exception as e:
+            raise RuntimeError(f"Failed to get branches for dataset {dataset_rid}: {e}")
+
+    def create_branch(
+        self, dataset_rid: str, branch_name: str, parent_branch: str = "master"
+    ) -> Dict[str, Any]:
+        """
+        Create a new branch for a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            branch_name: Name for the new branch
+            parent_branch: Parent branch to branch from
+
+        Returns:
+            Created branch information
+        """
+        try:
+            branch = self.service.create_branch(
+                dataset_rid=dataset_rid,
+                branch_name=branch_name,
+                parent_branch=parent_branch,
+            )
+
+            return {
+                "name": branch.name,
+                "dataset_rid": dataset_rid,
+                "parent_branch": parent_branch,
+                "transaction_rid": getattr(branch, "transaction_rid", None),
+            }
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to create branch '{branch_name}' for dataset {dataset_rid}: {e}"
+            )
+
+    def get_transactions(
+        self, dataset_rid: str, branch: str = "master"
+    ) -> List[Dict[str, Any]]:
+        """
+        Get list of transactions for a dataset branch.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            branch: Dataset branch name
+
+        Returns:
+            List of transaction information dictionaries
+        """
+        try:
+            # Note: This method may not be available in all SDK versions
+            transactions = self.service.list_transactions(
+                dataset_rid=dataset_rid, branch=branch
+            )
+
+            return [
+                {
+                    "transaction_rid": getattr(transaction, "rid", None),
+                    "created_time": getattr(transaction, "created_time", None),
+                    "created_by": getattr(transaction, "created_by", None),
+                    "status": getattr(transaction, "status", None),
+                }
+                for transaction in transactions
+            ]
+        except AttributeError:
+            # Method not available in this SDK version
+            raise NotImplementedError(
+                "Transaction listing is not supported by this SDK version. "
+                "This feature may require a newer version of foundry-platform-python."
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get transactions for dataset {dataset_rid}: {e}"
+            )
+
+    def get_views(self, dataset_rid: str) -> List[Dict[str, Any]]:
+        """
+        Get list of views for a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+
+        Returns:
+            List of view information dictionaries
+        """
+        try:
+            # Note: This method may not be available in all SDK versions
+            views = self.service.list_views(dataset_rid=dataset_rid)
+
+            return [
+                {
+                    "view_rid": getattr(view, "rid", None),
+                    "name": getattr(view, "name", None),
+                    "description": getattr(view, "description", None),
+                    "created_time": getattr(view, "created_time", None),
+                    "created_by": getattr(view, "created_by", None),
+                }
+                for view in views
+            ]
+        except AttributeError:
+            # Method not available in this SDK version
+            raise NotImplementedError(
+                "View listing is not supported by this SDK version. "
+                "This feature may require a newer version of foundry-platform-python."
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to get views for dataset {dataset_rid}: {e}")
+
+    def create_view(
+        self, dataset_rid: str, view_name: str, description: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Create a new view for a dataset.
+
+        Args:
+            dataset_rid: Dataset Resource Identifier
+            view_name: Name for the new view
+            description: Optional description for the view
+
+        Returns:
+            Created view information
+        """
+        try:
+            # Note: This method may not be available in all SDK versions
+            view = self.service.create_view(
+                dataset_rid=dataset_rid,
+                name=view_name,
+                description=description,
+            )
+
+            return {
+                "view_rid": getattr(view, "rid", None),
+                "name": view_name,
+                "description": description,
+                "dataset_rid": dataset_rid,
+                "created_time": getattr(view, "created_time", None),
+                "created_by": getattr(view, "created_by", None),
+            }
+        except AttributeError:
+            # Method not available in this SDK version
+            raise NotImplementedError(
+                "View creation is not supported by this SDK version. "
+                "This feature may require a newer version of foundry-platform-python."
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to create view '{view_name}' for dataset {dataset_rid}: {e}"
+            )
+
     def _format_dataset_info(self, dataset: Any) -> Dict[str, Any]:
         """
         Format dataset information for consistent output.
@@ -82,9 +430,19 @@ class DatasetService(BaseService):
         Returns:
             Formatted dataset information dictionary
         """
-        # The v2 Dataset object only has rid, name, and parent_folder_rid
+        # The v2 Dataset object has different attributes
         return {
-            "rid": dataset.rid,
-            "name": dataset.name,
-            "parent_folder_rid": dataset.parent_folder_rid,
+            "rid": getattr(dataset, "rid", "unknown"),
+            "name": getattr(dataset, "name", "Unknown"),
+            "description": getattr(dataset, "description", ""),
+            "path": getattr(dataset, "path", None),
+            "created": getattr(dataset, "created", None),
+            "modified": getattr(dataset, "modified", None),
+            # Try to get additional attributes that might exist
+            "created_time": getattr(dataset, "created_time", None),
+            "created_by": getattr(dataset, "created_by", None),
+            "last_modified": getattr(dataset, "last_modified", None),
+            "size_bytes": getattr(dataset, "size_bytes", None),
+            "schema_id": getattr(dataset, "schema_id", None),
+            "parent_folder_rid": getattr(dataset, "parent_folder_rid", None),
         }

pltr/services/folder.py

@@ -0,0 +1,167 @@
+"""
+Folder service wrapper for Foundry SDK filesystem API.
+"""
+
+from typing import Any, Optional, Dict, List
+
+from .base import BaseService
+
+
+class FolderService(BaseService):
+    """Service wrapper for Foundry folder operations using filesystem API."""
+
+    def _get_service(self) -> Any:
+        """Get the Foundry filesystem service."""
+        return self.client.filesystem
+
+    def create_folder(
+        self, display_name: str, parent_folder_rid: str
+    ) -> Dict[str, Any]:
+        """
+        Create a new folder.
+
+        Args:
+            display_name: Folder display name
+            parent_folder_rid: Parent folder RID (use 'ri.compass.main.folder.0' for root)
+
+        Returns:
+            Created folder information
+        """
+        try:
+            folder = self.service.Folder.create(
+                display_name=display_name,
+                parent_folder_rid=parent_folder_rid,
+                preview=True,
+            )
+            return self._format_folder_info(folder)
+        except Exception as e:
+            raise RuntimeError(f"Failed to create folder '{display_name}': {e}")
+
+    def get_folder(self, folder_rid: str) -> Dict[str, Any]:
+        """
+        Get information about a specific folder.
+
+        Args:
+            folder_rid: Folder Resource Identifier
+
+        Returns:
+            Folder information dictionary
+        """
+        try:
+            folder = self.service.Folder.get(folder_rid, preview=True)
+            return self._format_folder_info(folder)
+        except Exception as e:
+            raise RuntimeError(f"Failed to get folder {folder_rid}: {e}")
+
+    def list_children(
+        self,
+        folder_rid: str,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        List all child resources of a folder.
+
+        Args:
+            folder_rid: Folder Resource Identifier
+            page_size: Number of items per page (optional)
+            page_token: Pagination token (optional)
+
+        Returns:
+            List of child resources
+        """
+        try:
+            children = []
+            # The children method returns an iterator
+            for child in self.service.Folder.children(
+                folder_rid, page_size=page_size, page_token=page_token, preview=True
+            ):
+                children.append(self._format_resource_info(child))
+            return children
+        except Exception as e:
+            raise RuntimeError(f"Failed to list children of folder {folder_rid}: {e}")
+
+    def get_folders_batch(self, folder_rids: List[str]) -> List[Dict[str, Any]]:
+        """
+        Get multiple folders in a single request.
+
+        Args:
+            folder_rids: List of folder RIDs (max 1000)
+
+        Returns:
+            List of folder information dictionaries
+        """
+        if len(folder_rids) > 1000:
+            raise ValueError("Maximum batch size is 1000 folders")
+
+        try:
+            response = self.service.Folder.get_batch(body=folder_rids, preview=True)
+            folders = []
+            for folder in response.folders:
+                folders.append(self._format_folder_info(folder))
+            return folders
+        except Exception as e:
+            raise RuntimeError(f"Failed to get folders batch: {e}")
+
+    def _format_folder_info(self, folder: Any) -> Dict[str, Any]:
+        """
+        Format folder information for consistent output.
+
+        Args:
+            folder: Folder object from Foundry SDK
+
+        Returns:
+            Formatted folder information dictionary
+        """
+        return {
+            "rid": getattr(folder, "rid", None),
+            "display_name": getattr(folder, "display_name", None),
+            "description": getattr(folder, "description", None),
+            "created": self._format_timestamp(getattr(folder, "created", None)),
+            "modified": self._format_timestamp(getattr(folder, "modified", None)),
+            "parent_folder_rid": getattr(folder, "parent_folder_rid", None),
+            "type": "folder",
+        }
+
+    def _format_resource_info(self, resource: Any) -> Dict[str, Any]:
+        """
+        Format resource information (can be folder, dataset, etc.).
+
+        Args:
+            resource: Resource object from Foundry SDK
+
+        Returns:
+            Formatted resource information dictionary
+        """
+        resource_type = getattr(resource, "type", "unknown")
+        base_info = {
+            "rid": getattr(resource, "rid", None),
+            "display_name": getattr(resource, "display_name", None),
+            "type": resource_type,
+        }
+
+        # Add type-specific fields
+        if resource_type == "folder":
+            base_info["description"] = getattr(resource, "description", None)
+        elif resource_type == "dataset":
+            base_info["name"] = getattr(resource, "name", None)
+
+        return base_info
+
+    def _format_timestamp(self, timestamp: Any) -> Optional[str]:
+        """
+        Format timestamp for display.
+
+        Args:
+            timestamp: Timestamp object from SDK
+
+        Returns:
+            Formatted timestamp string or None
+        """
+        if timestamp is None:
+            return None
+
+        # Handle different timestamp formats from the SDK
+        if hasattr(timestamp, "time"):
+            return str(timestamp.time)
+        return str(timestamp)