manta-node 0.5b0.dev9__py3-none-any.whl

manta_node/infrastructure/filesystem/dataset_manager.py

@@ -0,0 +1,469 @@
+import os
+from logging import getLogger
+from pathlib import Path
+from typing import Any, Dict, Generator, List, Optional
+
+from manta_common.build.common.informations import DatasetOverviews, Overview
+from manta_common.const import CHUNK_SIZE
+from manta_common.errors import MantaSecurityError, MantaValidationError
+from manta_common.traces import Tracer
+
+__all__ = ["DatasetManager"]
+
+# File type to dtype mapping
+DTYPES = {
+    ".txt": "text",
+    ".csv": "csv",
+    ".npx": "numpy",
+    ".npy": "numpy",
+    ".npz": "numpy",
+    ".np": "numpy",
+    ".jpg": "image",
+    ".jpeg": "image",
+    ".png": "image",
+    ".tif": "image",
+    ".tiff": "image",
+    ".bin": "binary",
+}
+
+
+class DatasetManager:
+    """
+    Dataset Manager that handles parsing, retrieval, and metadata extraction
+    for datasets connected to a Node.
+
+    Parameters
+    ----------
+
+        Node ID
+    data_directories : Dict[str, str]
+        Dictionary mapping dataset names to folder paths
+
+    """
+
+    def __init__(self, data_directories: Optional[Dict[str, str]] = None):
+        self.tracer = Tracer(getLogger(__name__))
+        self.data_directories = {
+            name: Path(path) for name, path in (data_directories or {}).items()
+        }
+        # Parse the datasets to extract metadata
+        self._dataset_overviews = self._parse_datasets()
+
+        self.tracer.info(
+            f"Dataset manager initialized with {len(self.data_directories)} datasets"
+        )
+
+    def update_datasets(self, data_directories: Dict[str, str]) -> None:
+        """Update dataset mappings at runtime.
+
+        Parameters
+        ----------
+        data_directories : Dict[str, str]
+            New dictionary mapping dataset names to folder paths
+        """
+        self.tracer.info("Updating dataset mappings")
+
+        # Convert paths to Path objects
+        new_directories = {name: Path(path) for name, path in data_directories.items()}
+
+        # Check what's added and removed
+        added = set(new_directories.keys()) - set(self.data_directories.keys())
+        removed = set(self.data_directories.keys()) - set(new_directories.keys())
+
+        # Update the directories
+        self.data_directories = new_directories
+
+        # Re-parse datasets to get updated metadata
+        self._dataset_overviews = self._parse_datasets()
+
+        if added:
+            self.tracer.info(f"Added datasets: {added}")
+        if removed:
+            self.tracer.info(f"Removed datasets: {removed}")
+
+        self.tracer.info(
+            f"Dataset manager updated with {len(self.data_directories)} datasets"
+        )
+
+    def _parse_datasets(self) -> Dict[str, Overview]:
+        """
+        Parse all datasets and extract metadata
+
+        Returns
+        -------
+        Dict[str, Overview]
+            Dictionary mapping dataset names to their overviews
+        """
+        result = {}
+        for name, path in self.data_directories.items():
+            result[name] = self._parse_dataset(path)
+            self.tracer.info(f"Parsed dataset '{name}' with : {result[name]}")
+        return result
+
+    def _parse_dataset(self, folder_path: Path) -> Overview:
+        """
+        Parse a single dataset folder and extract metadata for all files
+
+        Parameters
+        ----------
+        folder_path : Path
+            Path to the dataset folder
+
+        Returns
+        -------
+        Overview
+            Dataset overview
+
+        Raises
+        ------
+        MantaValidationError
+            If the dataset path does not exist
+        """
+        if not folder_path.exists():
+            raise MantaValidationError(
+                f"Dataset path does not exist: {folder_path}. "
+                "Please create the dataset directory or file before starting the node."
+            )
+
+        if folder_path.is_file():
+            return Overview(
+                name=folder_path.stem,
+                description=f"File in {folder_path.name}",
+                dtype=DTYPES.get(folder_path.suffix.lower(), "binary"),
+            )
+        else:
+            return Overview(
+                name=folder_path.name,
+                description=f"Folder in {folder_path.name}",
+                dtype="folder",
+            )
+
+    def get_dataset_as_bytes(self, dataset_name: str) -> Generator[bytes, Any, None]:
+        """
+        Get binary data from a file
+
+        Parameters
+        ----------
+        dataset_name : str
+            Name of the dataset
+
+        Returns
+        -------
+        Generator[bytes, Any, None]
+            Generator yielding binary data from the file
+
+        Raises
+        ------
+        FileNotFoundError
+            If the file or dataset doesn't exist
+        ValueError
+            If the dataset is a folder
+        """
+        if dataset_name not in self._dataset_overviews:
+            raise FileNotFoundError(f"Dataset '{dataset_name}' not found")
+
+        if self._dataset_overviews[dataset_name].dtype == "folder":
+            raise ValueError(f"Dataset '{dataset_name}' is a folder")
+        elif self._dataset_overviews[dataset_name].dtype in [
+            "binary",
+            "csv",
+            "numpy",
+            "image",
+            "text",
+        ]:
+            return self.get_file_chunks(self.data_directories[dataset_name])
+
+        raise FileNotFoundError(f"Unknown dataset type for '{dataset_name}'")
+
+    def read_file_lines(
+        self,
+        dataset_name: str,
+        file_path: str,
+        encoding: str = "utf-8",
+        errors: str = "strict",
+        newline: str = "\n",
+    ) -> Generator[bytes, Any, None]:
+        """
+        Read the lines of a file
+
+        Parameters
+        ----------
+        dataset_name : str
+            Name of the dataset
+        file_path : str
+            Path to the file within the dataset
+        encoding : str, optional
+            Text encoding, by default "utf-8"
+        errors : str, optional
+            How to handle encoding errors, by default "strict"
+        newline : str, optional
+            Line ending character, by default "\n"
+
+        Returns
+        -------
+        Generator[bytes, Any, None]
+            Generator yielding lines as bytes
+
+        Raises
+        ------
+        FileNotFoundError
+            If the file or dataset doesn't exist
+        """
+        if not self.check_exists(dataset_name, file_path):
+            raise FileNotFoundError(f"File not found: {file_path}")
+        full_path = self.get_dataset_file_path(dataset_name, file_path)
+        with full_path.open(
+            mode="r", encoding=encoding, errors=errors, newline=newline
+        ) as f:
+            for line in f:
+                yield line.encode()
+
+    def get_file_chunks(self, file_path: Path) -> Generator[bytes, Any, None]:
+        """
+        Get binary data from a file in chunks
+
+        Parameters
+        ----------
+        file_path : Path
+            Path to the file within the dataset
+
+        Yields
+        ------
+        bytes
+            Chunks of binary data
+
+        Raises
+        ------
+        FileNotFoundError
+            If the file or dataset doesn't exist
+        """
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        with file_path.open("rb") as f:
+            while chunk := f.read(CHUNK_SIZE):
+                yield chunk
+
+    def _validate_file_path(self, file_path: str) -> str:
+        """
+        Validate and sanitize file paths to prevent directory traversal.
+
+        Parameters
+        ----------
+        file_path : str
+            The file path to validate
+
+        Returns
+        -------
+        str
+            The validated and sanitized file path
+
+        Raises
+        ------
+        MantaNodeError
+            If the path is invalid or contains directory traversal attempts
+        """
+        if not file_path:
+            return file_path
+
+        # Normalize the path to handle redundant separators and up-level references
+        normalized = os.path.normpath(file_path)
+
+        # Check for directory traversal attempts
+        if normalized.startswith("..") or os.path.isabs(normalized):
+            self.tracer.warning(
+                f"Security violation detected: directory traversal attempt in path: {file_path}"
+            )
+            raise MantaSecurityError(
+                f"Invalid file path - directory traversal attempt: {file_path}"
+            )
+
+        # Additional security: check for null bytes and other suspicious patterns
+        if "\x00" in normalized:
+            self.tracer.warning(
+                f"Security violation detected: null bytes in path: {file_path}"
+            )
+            raise MantaSecurityError(
+                f"Invalid file path - contains null bytes: {file_path}"
+            )
+
+        return normalized
+
+    def get_dataset_file_path(self, dataset_name: str, file_path: str) -> Path:
+        """
+        Get the full path to a dataset file with security validation
+
+        Parameters
+        ----------
+        dataset_name : str
+            Name of the dataset
+        file_path : str
+            Path to the file within the dataset
+
+        Returns
+        -------
+        Path
+            Full path to the file
+
+        Raises
+        ------
+        FileNotFoundError
+            If the file or dataset doesn't exist
+        MantaNodeError
+            If the path contains directory traversal attempts
+        """
+        if root := self.data_directories.get(dataset_name):
+            # Validate and sanitize the file path
+            validated_path = self._validate_file_path(file_path)
+
+            # Construct the full path
+            full_path = root / validated_path
+
+            # Ensure the resolved path is within the dataset directory
+            try:
+                # Resolve both paths to absolute paths to check containment
+                resolved_full = full_path.resolve()
+                resolved_root = root.resolve()
+
+                # Check if the resolved path is within the dataset directory
+                resolved_full.relative_to(resolved_root)
+            except ValueError:
+                raise MantaSecurityError(
+                    f"Path traversal attempt detected: {file_path} resolves outside dataset directory"
+                )
+
+            return full_path
+        raise FileNotFoundError(f"Dataset '{dataset_name}' not found")
+
+    def list_files(
+        self, dataset_name: str, folder_path: str = ""
+    ) -> List[Dict[str, Any]]:
+        """
+        List files in a dataset folder with security validation
+
+        Parameters
+        ----------
+        dataset_name : str
+            Name of the dataset
+        folder_path : str, optional
+            Path to the folder within the dataset, by default ""
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            List of file information including name and whether it's a file or folder
+
+        Raises
+        ------
+        FileNotFoundError
+            If the folder or dataset doesn't exist
+        MantaNodeError
+            If the path contains directory traversal attempts
+        """
+        if root := self.data_directories.get(dataset_name):
+            # Validate and sanitize the folder path
+            validated_folder_path = self._validate_file_path(folder_path)
+
+            # Construct the directory path
+            dir_path = root / validated_folder_path
+
+            # Ensure the resolved path is within the dataset directory
+            try:
+                resolved_dir = dir_path.resolve()
+                resolved_root = root.resolve()
+                resolved_dir.relative_to(resolved_root)
+            except ValueError:
+                raise MantaSecurityError(
+                    f"Path traversal attempt detected: {folder_path} resolves outside dataset directory"
+                )
+
+            if not dir_path.exists():
+                raise FileNotFoundError(f"Folder not found: {dir_path}")
+            return [
+                {"name": d.name, "is_file": d.is_file()} for d in dir_path.iterdir()
+            ]
+        raise FileNotFoundError(f"Dataset '{dataset_name}' not found")
+
+    def check_exists(self, dataset_name: str, file_path: str) -> bool:
+        """
+        Check if a file exists in a dataset with security validation
+
+        Parameters
+        ----------
+        dataset_name : str
+            Name of the dataset
+        file_path : str
+            Path to the file within the dataset
+
+        Returns
+        -------
+        bool
+            True if the file exists, False otherwise
+        """
+        try:
+            full_path = self.get_dataset_file_path(dataset_name, file_path)
+            return full_path.exists()
+        except (FileNotFoundError, MantaSecurityError, MantaValidationError):
+            return False
+
+    def to_proto(self) -> DatasetOverviews:
+        """
+        Convert all dataset overviews to a proto format
+
+        Returns
+        -------
+        DatasetOverviews
+            List of all dataset overviews
+        """
+        self.tracer.info(f"Dataset overviews: {self._dataset_overviews}")
+        return DatasetOverviews(
+            content=list(self._dataset_overviews.values()),
+        )
+
+    def validate_datasets(self) -> List[str]:
+        """
+        Validate that all configured dataset paths exist and are accessible.
+
+        Returns
+        -------
+        List[str]
+            List of validation issues (empty if all datasets are valid)
+        """
+        issues = []
+
+        for dataset_name, dataset_path in self.data_directories.items():
+            # Expand user paths (e.g., ~/)
+            expanded_path = dataset_path.expanduser()
+
+            # Check if the path exists
+            if not expanded_path.exists():
+                issues.append(
+                    f"Dataset '{dataset_name}' path does not exist: {dataset_path}"
+                )
+                continue
+
+            # Check if the path is readable (optimized using os.access)
+            try:
+                # Verify path type
+                if not (expanded_path.is_file() or expanded_path.is_dir()):
+                    issues.append(
+                        f"Dataset '{dataset_name}' path is neither a file nor a directory: {dataset_path}"
+                    )
+                    continue
+
+                # Use os.access for efficient permission check (no I/O overhead)
+                if not os.access(expanded_path, os.R_OK):
+                    issues.append(
+                        f"Dataset '{dataset_name}' is not readable (permission denied): {dataset_path}"
+                    )
+            except Exception as e:
+                issues.append(
+                    f"Dataset '{dataset_name}' access error: {dataset_path} - {e}"
+                )
+
+        return issues
+
+    def refresh(self):
+        """
+        Refresh the dataset metadata
+        """
+        self._dataset_overviews = self._parse_datasets()

manta_node/infrastructure/grpc/__init__.py

@@ -0,0 +1,11 @@
+"""gRPC infrastructure - Client and server implementations."""
+
+from .client import NodeClient
+from .local_servicer import LocalServicer
+from .world_servicer import WorldServicer
+
+__all__ = [
+    "NodeClient",
+    "LocalServicer",
+    "WorldServicer",
+]