pltr-cli 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

pltr/services/connectivity.py

@@ -58,6 +58,145 @@ class ConnectivityService(BaseService):
         except Exception as e:
             raise RuntimeError(f"Failed to get connection {connection_rid}: {e}")
 
+    def create_connection(
+        self,
+        display_name: str,
+        parent_folder_rid: str,
+        configuration: Dict[str, Any],
+        worker: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """
+        Create a new connection.
+
+        Args:
+            display_name: Display name for the connection
+            parent_folder_rid: Parent folder Resource Identifier
+            configuration: Connection configuration dictionary
+            worker: Worker configuration dictionary
+
+        Returns:
+            Created connection information dictionary
+        """
+        try:
+            connection = self.connections_service.Connection.create(
+                configuration=configuration,
+                display_name=display_name,
+                parent_folder_rid=parent_folder_rid,
+                worker=worker,
+            )
+            return self._format_connection_info(connection)
+        except Exception as e:
+            raise RuntimeError(f"Failed to create connection '{display_name}': {e}")
+
+    def get_connection_configuration(self, connection_rid: str) -> Dict[str, Any]:
+        """
+        Get connection configuration.
+
+        Args:
+            connection_rid: Connection Resource Identifier
+
+        Returns:
+            Connection configuration dictionary
+        """
+        try:
+            config = self.connections_service.Connection.get_configuration(
+                connection_rid
+            )
+            return {"connection_rid": connection_rid, "configuration": config}
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get configuration for connection {connection_rid}: {e}"
+            )
+
+    def update_export_settings(
+        self, connection_rid: str, export_settings: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Update connection export settings.
+
+        Args:
+            connection_rid: Connection Resource Identifier
+            export_settings: Export settings dictionary
+
+        Returns:
+            Status dictionary
+        """
+        try:
+            self.connections_service.Connection.update_export_settings(
+                connection_rid=connection_rid,
+                export_settings=export_settings,
+            )
+            return {
+                "connection_rid": connection_rid,
+                "status": "export settings updated",
+            }
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to update export settings for connection {connection_rid}: {e}"
+            )
+
+    def update_secrets(
+        self, connection_rid: str, secrets: Dict[str, str]
+    ) -> Dict[str, Any]:
+        """
+        Update connection secrets.
+
+        Args:
+            connection_rid: Connection Resource Identifier
+            secrets: Dictionary mapping secret names to values
+
+        Returns:
+            Status dictionary
+        """
+        try:
+            self.connections_service.Connection.update_secrets(
+                connection_rid=connection_rid,
+                secrets=secrets,
+            )
+            return {"connection_rid": connection_rid, "status": "secrets updated"}
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to update secrets for connection {connection_rid}: {e}"
+            )
+
+    def upload_custom_jdbc_drivers(
+        self, connection_rid: str, file_path: str
+    ) -> Dict[str, Any]:
+        """
+        Upload custom JDBC drivers to a connection.
+
+        Args:
+            connection_rid: Connection Resource Identifier
+            file_path: Path to the JAR file
+
+        Returns:
+            Updated connection information dictionary
+        """
+        from pathlib import Path
+
+        file_path_obj = Path(file_path)
+
+        if not file_path_obj.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+
+        if not file_path_obj.suffix.lower() == ".jar":
+            raise ValueError(f"File must be a JAR file: {file_path}")
+
+        try:
+            with open(file_path_obj, "rb") as f:
+                file_content = f.read()
+
+            connection = self.connections_service.Connection.upload_custom_jdbc_drivers(
+                connection_rid=connection_rid,
+                body=file_content,
+                file_name=file_path_obj.name,
+            )
+            return self._format_connection_info(connection)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to upload JDBC driver to connection {connection_rid}: {e}"
+            )
+
     def create_file_import(
         self,
         connection_rid: str,

pltr/services/copy.py

@@ -0,0 +1,391 @@
+"""
+High-level service that copies Foundry resources between Compass folders.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path, PurePosixPath
+from typing import Dict, List, Optional
+import tempfile
+import traceback
+
+from rich.console import Console
+
+from .dataset import DatasetService
+from .folder import FolderService
+from .resource import ResourceService
+
+
+# Known resource types for exact matching (avoids false positives with substring matching)
+DATASET_TYPES = frozenset({"foundry_dataset", "dataset"})
+FOLDER_TYPES = frozenset({"folder", "compass_folder"})
+
+
+@dataclass
+class CopyStats:
+    """Summary of a copy operation."""
+
+    folders: int = 0
+    datasets: int = 0
+    skipped: int = 0
+    errors: int = 0
+
+    def as_dict(self) -> Dict[str, int]:
+        """Return summary as a dictionary."""
+        return {
+            "folders_copied": self.folders,
+            "datasets_copied": self.datasets,
+            "skipped": self.skipped,
+            "errors": self.errors,
+        }
+
+
+class CopyService:
+    """Copy datasets or folders (and their children) into another Compass folder."""
+
+    def __init__(
+        self,
+        *,
+        profile: Optional[str] = None,
+        branch: str = "master",
+        name_suffix: str = "-copy",
+        copy_schema: bool = True,
+        dry_run: bool = False,
+        debug: bool = False,
+        fail_fast: bool = False,
+        console: Optional[Console] = None,
+    ):
+        self.profile = profile
+        self.branch = branch
+        self.name_suffix = name_suffix
+        self.copy_schema = copy_schema
+        self.dry_run = dry_run
+        self.debug = debug
+        self.fail_fast = fail_fast
+        self.console = console or Console()
+
+        self.dataset_service = DatasetService(profile=profile)
+        self.folder_service = FolderService(profile=profile)
+        self.resource_service = ResourceService(profile=profile)
+
+        self.stats = CopyStats()
+        self._skipped_messages: List[str] = []
+
+    # ------------------------------------------------------------------ public
+    def copy_resource(
+        self, source_rid: str, target_folder_rid: str, recursive: bool = False
+    ) -> Dict[str, int]:
+        """
+        Copy a resource by RID into a destination folder.
+
+        Args:
+            source_rid: Resource Identifier to copy.
+            target_folder_rid: Destination Compass folder RID.
+            recursive: Required when copying folders. Recursively copies contents.
+
+        Returns:
+            Dictionary summary of copied items.
+        """
+        self._reset_stats()
+
+        resource = self.resource_service.get_resource(source_rid)
+        resource_type = (resource.get("type") or "").lower()
+
+        if resource_type in DATASET_TYPES:
+            self._log_info(
+                f"Copying dataset '{self._get_resource_name(resource)}' "
+                f"({source_rid}) → folder {target_folder_rid}"
+            )
+            try:
+                self._copy_dataset(resource, target_folder_rid)
+            except Exception:
+                self.stats.errors += 1
+                raise
+        elif resource_type in FOLDER_TYPES:
+            if not recursive:
+                raise RuntimeError(
+                    "Source resource is a folder. Pass --recursive to copy folder contents."
+                )
+            try:
+                self._copy_folder(resource, target_folder_rid)
+            except Exception:
+                self.stats.errors += 1
+                raise
+        else:
+            raise RuntimeError(
+                f"Copy is only supported for datasets and folders. Resource "
+                f"{source_rid} is of type '{resource.get('type')}'."
+            )
+
+        self._print_summary()
+        return self.stats.as_dict()
+
+    # ----------------------------------------------------------------- datasets
+    def _copy_dataset(
+        self, dataset_info: Dict[str, str], target_folder_rid: str
+    ) -> None:
+        dataset_rid = dataset_info["rid"]
+        dataset_name = self._get_resource_name(dataset_info)
+        new_name = self._derive_name(dataset_name)
+
+        if self.dry_run:
+            self._log_warning(
+                f"[DRY-RUN] Would copy dataset '{dataset_name}' ({dataset_rid}) "
+                f"→ '{new_name}' in folder {target_folder_rid}"
+            )
+            self.stats.datasets += 1
+            return
+
+        new_dataset = self.dataset_service.create_dataset(new_name, target_folder_rid)
+        new_dataset_rid = new_dataset["rid"]
+        self._log_success(
+            f"Created dataset '{new_name}' ({new_dataset_rid}) in {target_folder_rid}"
+        )
+
+        try:
+            if self.copy_schema:
+                self._copy_dataset_schema(dataset_rid, new_dataset_rid)
+
+            self._copy_dataset_files(dataset_rid, new_dataset_rid)
+            self.stats.datasets += 1
+            self._log_success(f"Finished copying dataset to {new_dataset_rid}")
+        except Exception as exc:
+            # Clean up partially created dataset on failure
+            self._log_warning(
+                f"  Deleting partially created dataset {new_dataset_rid} due to error"
+            )
+            try:
+                self.dataset_service.delete_dataset(new_dataset_rid)
+                self._log_info(f"  Deleted dataset {new_dataset_rid}")
+            except Exception as delete_exc:
+                self._log_error(
+                    f"  Failed to delete dataset {new_dataset_rid}: {delete_exc}"
+                )
+            raise exc
+
+    def _copy_dataset_schema(self, source_rid: str, target_rid: str) -> None:
+        """Copy schema metadata, warning if not available."""
+        try:
+            schema_info = self.dataset_service.get_schema(source_rid)
+            schema = schema_info.get("schema")
+            if schema:
+                self.dataset_service.put_schema(target_rid, schema)
+                self._log_info("  Copied dataset schema")
+            else:
+                self._log_info("  Source dataset has no schema to copy")
+        except Exception as exc:
+            self._log_warning(f"  Could not copy schema for {source_rid}: {exc}")
+            if self.debug:
+                traceback.print_exc()
+
+    def _copy_dataset_files(self, source_rid: str, target_rid: str) -> None:
+        """Download files from source and upload them into the target dataset."""
+        files = self.dataset_service.list_files(source_rid, branch=self.branch)
+        if not files:
+            self._log_info(
+                "  Source dataset does not expose any files on the requested branch"
+            )
+            return
+
+        # Wrap transaction creation in try block to ensure cleanup on any failure
+        transaction_rid = None
+        transaction_created = False
+        try:
+            transaction = self.dataset_service.create_transaction(
+                target_rid, branch=self.branch, transaction_type="SNAPSHOT"
+            )
+            # Check both possible key names for the transaction RID
+            transaction_rid = transaction.get("transaction_rid") or transaction.get(
+                "rid"
+            )
+            if not transaction_rid:
+                raise RuntimeError(
+                    f"Transaction response missing RID: {list(transaction.keys())}"
+                )
+            transaction_created = True
+
+            with tempfile.TemporaryDirectory(prefix="pltr-cp-") as tmpdir:
+                temp_dir = Path(tmpdir)
+                for file_info in files:
+                    dataset_path = file_info.get("path")
+                    if not dataset_path:
+                        self._report_skip(
+                            {
+                                "rid": source_rid,
+                                "display_name": dataset_path or "unknown file",
+                            },
+                            "File path missing; skipping file.",
+                        )
+                        continue
+
+                    local_rel_path = self._sanitize_local_path(dataset_path)
+                    download_path = temp_dir / local_rel_path
+                    download_path.parent.mkdir(parents=True, exist_ok=True)
+
+                    self.dataset_service.download_file(
+                        source_rid, dataset_path, download_path, branch=self.branch
+                    )
+                    self._upload_dataset_file(
+                        target_rid, dataset_path, download_path, transaction_rid
+                    )
+                    self._log_info(f"  Copied dataset file: {dataset_path}")
+
+            self.dataset_service.commit_transaction(target_rid, transaction_rid)
+            self._log_info(f"  Committed transaction {transaction_rid}")
+        except Exception as exc:
+            # Only attempt rollback if we successfully created a transaction
+            if transaction_created and transaction_rid:
+                try:
+                    self.dataset_service.abort_transaction(target_rid, transaction_rid)
+                    self._log_warning(f"  Rolled back transaction {transaction_rid}")
+                except Exception:
+                    self._log_warning(
+                        f"  Failed to roll back transaction {transaction_rid}"
+                    )
+
+            if self.debug:
+                traceback.print_exc()
+
+            raise RuntimeError(
+                f"{type(exc).__name__} while copying files from {source_rid} to {target_rid}: {exc}"
+            ) from exc
+
+    def _upload_dataset_file(
+        self,
+        dataset_rid: str,
+        dataset_path: str,
+        local_file: Path,
+        transaction_rid: str,
+    ) -> None:
+        """Upload the downloaded file bytes back into a dataset."""
+        with open(local_file, "rb") as handle:
+            body = handle.read()
+
+        remote_path = PurePosixPath(dataset_path.lstrip("/")).as_posix()
+        upload_kwargs = {
+            "dataset_rid": dataset_rid,
+            "file_path": remote_path,
+            "body": body,
+            "transaction_rid": transaction_rid,
+        }
+        # API only allows either branchName or transactionRid. When using a transaction
+        # we must omit branchName to avoid InvalidParameterCombination.
+        if not transaction_rid:
+            upload_kwargs["branch_name"] = self.branch
+
+        self.dataset_service.service.Dataset.File.upload(**upload_kwargs)
+
+    # ------------------------------------------------------------------ folders
+    def _copy_folder(self, folder_info: Dict[str, str], target_folder_rid: str) -> None:
+        folder_rid = folder_info["rid"]
+        folder_name = self._get_resource_name(folder_info)
+        new_name = self._derive_name(folder_name)
+
+        if self.dry_run:
+            self._log_warning(
+                f"[DRY-RUN] Would copy folder '{folder_name}' ({folder_rid}) "
+                f"→ '{new_name}' in folder {target_folder_rid}"
+            )
+            new_folder_rid = "<dry-run>"
+        else:
+            created = self.folder_service.create_folder(new_name, target_folder_rid)
+            new_folder_rid = created["rid"]
+            self._log_success(
+                f"Created folder '{new_name}' ({new_folder_rid}) in {target_folder_rid}"
+            )
+
+        self.stats.folders += 1
+
+        children = self.folder_service.list_children(folder_rid)
+        if not children:
+            self._log_info("  Folder has no children to copy.")
+            return
+
+        for child in children:
+            resource_type_raw = child.get("type") or ""
+            resource_type = resource_type_raw.lower()
+
+            try:
+                if resource_type in FOLDER_TYPES:
+                    self._copy_folder(child, new_folder_rid)
+                elif resource_type in DATASET_TYPES:
+                    self._copy_dataset(child, new_folder_rid)
+                else:
+                    self._report_skip(
+                        child, f"Unsupported child type '{resource_type_raw}'"
+                    )
+            except Exception as exc:
+                self.stats.errors += 1
+                self._log_error(
+                    f"Failed to copy child {child.get('rid', 'unknown RID')}: {exc}"
+                )
+                if self.debug:
+                    traceback.print_exc()
+                if self.fail_fast:
+                    raise RuntimeError(f"Stopping due to --fail-fast: {exc}") from exc
+
+    # ------------------------------------------------------------------ helpers
+    def _get_resource_name(self, resource: Dict[str, str]) -> str:
+        """Extract display name from a resource, with fallbacks."""
+        return (
+            resource.get("display_name")
+            or resource.get("name")
+            or resource.get("rid", "unknown")
+        )
+
+    def _sanitize_local_path(self, original_path: str) -> str:
+        """Return a safe relative path for storing temporary files."""
+        # Check for path traversal BEFORE normalization to prevent attacks
+        if ".." in original_path:
+            raise ValueError(f"Path traversal detected: {original_path}")
+        clean = PurePosixPath(original_path.lstrip("/"))
+        # Verify no parent directory references remain after normalization
+        if ".." in clean.parts:
+            raise ValueError(f"Invalid path after normalization: {original_path}")
+        return clean.as_posix()
+
+    def _derive_name(self, base_name: str) -> str:
+        """Append the configured suffix if provided."""
+        suffix = self.name_suffix or ""
+        if not suffix:
+            return base_name
+        if base_name.endswith(suffix):
+            return base_name
+        return f"{base_name}{suffix}"
+
+    def _report_skip(self, resource: Dict[str, str], reason: str) -> None:
+        """Track skipped resources."""
+        rid = resource.get("rid", "unknown RID")
+        display_name = resource.get("display_name") or resource.get("name") or rid
+        message = f"[SKIP] {display_name} ({rid}): {reason}"
+        self.console.print(message)
+        self.stats.skipped += 1
+        self._skipped_messages.append(message)
+
+    def _reset_stats(self) -> None:
+        self.stats = CopyStats()
+        self._skipped_messages = []
+
+    def _print_summary(self) -> None:
+        """Print summary after an operation."""
+        self.console.print(
+            "\nSummary: "
+            f"{self.stats.folders} folders copied, "
+            f"{self.stats.datasets} datasets copied, "
+            f"{self.stats.skipped} resources skipped, "
+            f"{self.stats.errors} errors"
+        )
+
+    # ------------------------------------------------------------------ logging
+    def _log_info(self, message: str) -> None:
+        self.console.print(message)
+
+    def _log_warning(self, message: str) -> None:
+        self.console.print(f"[yellow]{message}[/yellow]")
+
+    def _log_error(self, message: str) -> None:
+        self.console.print(f"[red]{message}[/red]")
+
+    def _log_success(self, message: str) -> None:
+        self.console.print(f"[green]{message}[/green]")

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
 
 
@@ -16,8 +18,6 @@ class DatasetService(BaseService):
         """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.
@@ -252,8 +252,9 @@ class DatasetService(BaseService):
             # Clean column name (remove special characters for field name)
             clean_name = col.strip().replace(" ", "_").replace("-", "_")
 
+            # SDK 1.69.0 expects FieldType enum but accepts strings at runtime
             fields.append(
-                DatasetFieldSchema(name=clean_name, type=field_type, nullable=nullable)
+                DatasetFieldSchema(name=clean_name, type=field_type, nullable=nullable)  # type: ignore[arg-type]
             )
 
         return DatasetSchema(field_schema_list=fields)
@@ -312,6 +313,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.
@@ -516,6 +540,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
@@ -540,6 +566,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/folder.py

@@ -4,6 +4,8 @@ Folder service wrapper for Foundry SDK filesystem API.
 
 from typing import Any, Optional, Dict, List
 
+from foundry_sdk.v2.filesystem.models import GetFoldersBatchRequestElement
+
 from .base import BaseService
 
 
@@ -95,7 +97,10 @@ class FolderService(BaseService):
             raise ValueError("Maximum batch size is 1000 folders")
 
         try:
-            response = self.service.Folder.get_batch(body=folder_rids, preview=True)
+            elements = [
+                GetFoldersBatchRequestElement(folder_rid=rid) for rid in folder_rids
+            ]
+            response = self.service.Folder.get_batch(body=elements, preview=True)
             folders = []
             for folder in response.folders:
                 folders.append(self._format_folder_info(folder))