streamlit-octostar-utils 0.4.1__tar.gz → 0.4.2.dev2__tar.gz

{streamlit_octostar_utils-0.4.1 → streamlit_octostar_utils-0.4.2.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.4.1
+Version: 0.4.2.dev2
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.4.1 → streamlit_octostar_utils-0.4.2.dev2}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.4.1"
+version = "0.4.2-dev.2"
 description = ""
 license = "MIT"
 authors = ["Octostar"]

streamlit_octostar_utils-0.4.2.dev2/streamlit_octostar_utils/api_crafter/contents.py

@@ -0,0 +1,914 @@
+"""
+Contents abstraction for NiFi entities.
+
+Provides a FileIO-like interface for handling entity contents with support for:
+- Full io.BufferedIOBase compatibility (read, write, seek, tell, etc.)
+- Seeking and streaming with HTTP Range requests
+- Multiple storage backends (memory, workspace attachments)
+- Serialization/deserialization via from_locator()/to_locator() methods
+- Lazy loading and efficient partial reads for large files
+
+Storage Backends:
+- MemoryContents: In-memory buffer (like BytesIO)
+- WorkspaceAttachmentContents: Octostar workspace attachments with HTTP Range support
+- TemporaryAttachmentContents: Octostar temporary blob storage (user temp bucket)
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, Any, Union, BinaryIO, List
+from enum import Enum
+import base64
+import httpx
+import logging
+from io import BytesIO, SEEK_SET, SEEK_CUR, SEEK_END
+from copy import deepcopy
+
+_logger = logging.getLogger(__name__)
+
+
+class ContentsLocation(Enum):
+    """Enumeration of supported content storage locations."""
+    MEMORY = "memory"
+    WORKSPACE_ATTACHMENT = "workspace_attachment"
+    TEMPORARY_ATTACHMENT = "temporary_attachment"
+
+
+class Contents(ABC):
+    """
+    Abstract base class for entity contents.
+    
+    Provides a FileIO-like interface with support for seeking, streaming,
+    and multiple storage backends. Each implementation manages its own
+    locator dictionary for serialization/deserialization.
+    """
+    
+    def __init__(
+        self,
+        entity_type: Optional[str] = None,
+        filetype: Optional[str] = None,
+        **kwargs
+    ):
+        self._entity_type = entity_type
+        self._filetype = filetype
+        self._closed = False
+    
+    # ==================== FileIO Interface ====================
+    
+    @abstractmethod
+    def read(self, size: int = -1) -> bytes:
+        """
+        Read and return up to size bytes, or all bytes if size is -1.
+        
+        Args:
+            size: Number of bytes to read. -1 means read all.
+            
+        Returns:
+            Bytes read from the stream.
+        """
+        pass
+    
+    @abstractmethod
+    def write(self, b: bytes) -> int:
+        """
+        Write bytes to the stream.
+        
+        Args:
+            b: Bytes to write.
+            
+        Returns:
+            Number of bytes written.
+        """
+        pass
+    
+    @abstractmethod
+    def seek(self, offset: int, whence: int = SEEK_SET) -> int:
+        """
+        Change stream position.
+        
+        Args:
+            offset: Offset relative to whence.
+            whence: SEEK_SET (0) = from start, SEEK_CUR (1) = from current, SEEK_END (2) = from end.
+            
+        Returns:
+            New absolute position.
+        """
+        pass
+    
+    @abstractmethod
+    def tell(self) -> int:
+        """
+        Return current stream position.
+        
+        Returns:
+            Current position in bytes.
+        """
+        pass
+    
+    def readable(self) -> bool:
+        return True
+    
+    def writable(self) -> bool:
+        return True
+    
+    def seekable(self) -> bool:
+        return True
+    
+    @abstractmethod
+    def flush(self):
+        """Flush write buffers."""
+        pass
+    
+    @abstractmethod
+    def close(self):
+        """Close the stream and release resources."""
+        self._closed = True
+    
+    @abstractmethod
+    def delete(self):
+        """Delete the contents from their storage backend."""
+        pass
+    
+    @property
+    def closed(self) -> bool:
+        """Check if stream is closed."""
+        return self._closed
+    
+    @abstractmethod
+    def truncate(self, size: Optional[int] = None) -> int:
+        """
+        Resize the stream to the given size.
+        
+        Args:
+            size: New size in bytes. If None, use current position.
+            
+        Returns:
+            New size.
+        """
+        pass
+    
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()
+        return False
+    
+    def readline(self, size: int = -1) -> bytes:
+        """
+        Read and return one line from the stream.
+        
+        Args:
+            size: Maximum number of bytes to read. -1 means no limit.
+            
+        Returns:
+            Bytes up to and including the newline character, or until EOF.
+        """
+        result = bytearray()
+        while True:
+            if size >= 0 and len(result) >= size:
+                break
+            byte = self.read(1)
+            if not byte:
+                break
+            result.extend(byte)
+            if byte == b'\n':
+                break
+        return bytes(result)
+    
+    def readlines(self, hint: int = -1) -> List[bytes]:
+        """
+        Read and return a list of lines from the stream.
+        
+        Args:
+            hint: Optional size hint. If total size of lines exceeds hint, no more lines are read.
+            
+        Returns:
+            List of lines.
+        """
+        lines = []
+        total_size = 0
+        while True:
+            line = self.readline()
+            if not line:
+                break
+            lines.append(line)
+            total_size += len(line)
+            if hint >= 0 and total_size >= hint:
+                break
+        return lines
+    
+    def writelines(self, lines: List[bytes]):
+        """
+        Write a list of lines to the stream.
+        
+        Args:
+            lines: List of byte strings to write.
+        """
+        for line in lines:
+            self.write(line)
+    
+    def readinto(self, b: bytearray) -> int:
+        """
+        Read bytes into a pre-allocated buffer.
+        
+        Args:
+            b: Pre-allocated bytearray to read into.
+            
+        Returns:
+            Number of bytes read.
+        """
+        data = self.read(len(b))
+        n = len(data)
+        b[:n] = data
+        return n
+    
+    def read1(self, size: int = -1) -> bytes:
+        """
+        Read up to size bytes with at most one read() call to the underlying stream.
+        
+        For most implementations, this is the same as read(). Subclasses may override
+        for optimization.
+        
+        Args:
+            size: Number of bytes to read. -1 means read all available.
+            
+        Returns:
+            Bytes read from the stream.
+        """
+        return self.read(size)
+    
+    # ==================== Locator Interface ====================
+    
+    @abstractmethod
+    def to_locator(self) -> Dict[str, Any]:
+        """
+        Serialize contents to a locator dictionary.
+        
+        The locator contains:
+        - location: The storage backend type
+        - Additional backend-specific parameters
+        
+        Returns:
+            Dictionary describing how to locate/access this content.
+        """
+        pass
+    
+    @staticmethod
+    def from_locator(locator: Optional[Dict[str, Any]], client=None) -> Optional['Contents']:
+        """
+        Deserialize contents from a locator dictionary.
+        
+        Factory method that creates the appropriate Contents subclass based on the
+        locator's "location" field. This replaces the old base64.b64decode logic in NiFi.
+        
+        Args:
+            locator: Locator dictionary (e.g., from request.contents_pointer)
+            client: Optional Octostar client for remote operations
+            
+        Returns:
+            Contents instance or None if locator is None/empty.
+            
+        Raises:
+            ValueError: If the location type is unknown
+        """
+        if not locator:
+            return None
+        
+        location = locator.get("location")
+        
+        match location:
+            case ContentsLocation.MEMORY.value:
+                return MemoryContents._from_locator(locator)
+            case ContentsLocation.WORKSPACE_ATTACHMENT.value:
+                return WorkspaceAttachmentContents._from_locator(locator, client)
+            case ContentsLocation.TEMPORARY_ATTACHMENT.value:
+                return TemporaryAttachmentContents._from_locator(locator, client)
+            case _:
+                raise ValueError(f"Unknown contents location type: {location}")
+    
+    # ==================== Utility Methods ====================
+    
+    def read_all(self) -> bytes:
+        """Read all contents and return as bytes."""
+        current_pos = self.tell()
+        self.seek(0, SEEK_SET)
+        data = self.read()
+        self.seek(current_pos, SEEK_SET)
+        return data
+    
+    def getvalue(self) -> bytes:
+        """
+        Return the entire contents without moving the position.
+        
+        This matches BytesIO.getvalue() behavior.
+        
+        Returns:
+            Entire contents as bytes.
+        """
+        return self.read_all()
+    
+    def write_all(self, data: bytes):
+        """Write all data, replacing existing contents."""
+        self.seek(0, SEEK_SET)
+        self.truncate(0)
+        self.write(data)
+        self.flush()
+    
+    def __len__(self) -> int:
+        """
+        Return the length of the contents.
+        
+        Returns:
+            Total size in bytes.
+        """
+        current_pos = self.tell()
+        self.seek(0, SEEK_END)
+        length = self.tell()
+        self.seek(current_pos, SEEK_SET)
+        return length
+
+
+class MemoryContents(Contents):
+    """In-memory contents implementation using BytesIO."""
+    
+    def __init__(
+        self,
+        entity_type: Optional[str] = None,
+        filetype: Optional[str] = None,
+        *,
+        initial_data: Optional[bytes] = None,
+        **kwargs
+    ):
+        super().__init__(entity_type, filetype, **kwargs)
+        self._buffer = BytesIO(initial_data or b"")
+    
+    def read(self, size: int = -1) -> bytes:
+        return self._buffer.read(size)
+    
+    def write(self, b: bytes) -> int:
+        return self._buffer.write(b)
+    
+    def seek(self, offset: int, whence: int = SEEK_SET) -> int:
+        return self._buffer.seek(offset, whence)
+    
+    def tell(self) -> int:
+        return self._buffer.tell()
+    
+    def flush(self):
+        self._buffer.flush()
+    
+    def close(self):
+        if not self._closed:
+            self._buffer.close()
+            super().close()
+    
+    def delete(self):
+        """Clear the in-memory buffer."""
+        self._buffer = BytesIO(b"")
+    
+    def truncate(self, size: Optional[int] = None) -> int:
+        return self._buffer.truncate(size)
+    
+    def getvalue(self) -> bytes:
+        return self._buffer.getvalue()
+    
+    def to_locator(self) -> Dict[str, Any]:
+        """
+        Serialize to locator with base64-encoded data.
+        
+        Returns:
+            {"location": "memory", "data": "<base64>", "entity_type": "...", "filetype": "..."}
+        """
+        data = self._buffer.getvalue()
+        locator = {
+            "location": ContentsLocation.MEMORY.value,
+            "data": base64.b64encode(data).decode('utf-8') if data else None
+        }
+        if self._entity_type:
+            locator["entity_type"] = self._entity_type
+        if self._filetype:
+            locator["filetype"] = self._filetype
+        return locator
+    
+    @staticmethod
+    def _from_locator(locator: Dict[str, Any]) -> 'MemoryContents':
+        """
+        Create MemoryContents from a locator dictionary.
+        
+        Args:
+            locator: Locator dictionary with base64-encoded data
+            
+        Returns:
+            New MemoryContents instance
+        """
+        data = locator.get("data")
+        initial_data = base64.b64decode(data) if data else None
+        return MemoryContents(
+            entity_type=locator.get("entity_type"),
+            filetype=locator.get("filetype"),
+            initial_data=initial_data
+        )
+
+
+class WorkspaceAttachmentContents(Contents):
+    """
+    Contents implementation for Octostar workspace attachments.
+    
+    Uses octostar-api utilities (read_attachment, get_attachment_url) to:
+    - Lazy load from workspace storage via read_attachment()
+    - Stream efficiently with HTTP Range requests (no need to download entire file)
+    - Support true seeking with configurable chunk sizes
+    - Buffer modifications in memory (write back handled at NiFi entity level)
+    
+    Presigned URL Handling:
+    - URLs are obtained via get_attachment_url() and cached
+    - On 403 (Forbidden) responses, URLs are refreshed and requests retried
+    - Follows the same retry pattern as octostar-api read_file utilities
+    """
+    
+    DEFAULT_CHUNK_SIZE = 8192
+    DEFAULT_URL_TIMEOUT = 120
+    
+    def __init__(
+        self,
+        entity_type: Optional[str] = None,
+        filetype: Optional[str] = None,
+        *,
+        workspace_id: str,
+        entity_id: str,
+        client,
+        initial_data: Optional[bytes] = None,
+        chunk_size: int = DEFAULT_CHUNK_SIZE,
+        **kwargs
+    ):
+        super().__init__(entity_type, filetype, **kwargs)
+        self._workspace_id = workspace_id
+        self._entity_id = entity_id
+        self._client = client
+        self._chunk_size = chunk_size
+        
+        self._buffer: Optional[BytesIO] = None
+        self._fully_loaded = False
+        self._modified = False
+        self._position = 0
+        self._size: Optional[int] = None
+        self._presigned_url: Optional[str] = None
+        self._http_client: Optional[httpx.Client] = None
+        
+        if initial_data is not None:
+            self._buffer = BytesIO(initial_data)
+            self._fully_loaded = True
+            self._size = len(initial_data)
+    
+    def _get_presigned_url(self) -> str:
+        """Get or refresh the presigned URL for the attachment."""
+        from octostar.utils.workspace import get_attachment_url
+        
+        return get_attachment_url.sync(
+            os_workspace=self._workspace_id,
+            os_entity_uid=self._entity_id,
+            client=self._client
+        )
+    
+    def _ensure_http_client(self):
+        """Ensure HTTP client is initialized."""
+        if not self._http_client:
+            self._http_client = httpx.Client(timeout=self.DEFAULT_URL_TIMEOUT)
+    
+    def _fetch_size(self) -> int:
+        """
+        Fetch the size of the remote attachment using HEAD request.
+        
+        Handles presigned URL expiration with automatic refresh and retry.
+        """
+        if self._size is not None:
+            return self._size
+        
+        if not self._presigned_url:
+            self._presigned_url = self._get_presigned_url()
+        
+        self._ensure_http_client()
+        max_retries = 3
+        
+        for attempt in range(max_retries):
+            response = self._http_client.head(self._presigned_url)
+            if response.status_code == 200:
+                self._size = int(response.headers.get('content-length', 0))
+                return self._size
+            elif response.status_code == 403:
+                _logger.debug(
+                    f"Presigned URL expired while fetching size, refreshing... (attempt {attempt + 1}/{max_retries})"
+                )
+                self._presigned_url = self._get_presigned_url()
+                continue
+            else:
+                response.raise_for_status()
+
+        raise ConnectionError(
+            f"Failed to fetch file size after {max_retries} attempts (URL kept expiring)"
+        )
+    
+    def _read_range(self, start: int, end: int) -> bytes:
+        """
+        Read a specific byte range using HTTP Range request.
+        
+        Handles presigned URL expiration with automatic refresh and retry.
+        
+        Args:
+            start: Start byte (inclusive)
+            end: End byte (inclusive)
+            
+        Returns:
+            Bytes from the specified range
+        """
+        if not self._presigned_url:
+            self._presigned_url = self._get_presigned_url()
+        
+        self._ensure_http_client()
+        
+        headers = {"Range": f"bytes={start}-{end}"}
+        max_retries = 3
+        
+        for attempt in range(max_retries):
+            response = self._http_client.get(self._presigned_url, headers=headers)
+            
+            if response.status_code == 206:
+                return response.content
+            elif response.status_code == 416:
+                return b""
+            elif response.status_code == 403:
+                _logger.debug(
+                    f"Presigned URL expired at byte {start}, refreshing... (attempt {attempt + 1}/{max_retries})"
+                )
+                self._presigned_url = self._get_presigned_url()
+                continue
+            else:
+                response.raise_for_status()
+                return response.content
+                
+        raise ConnectionError(
+            f"Failed to read range {start}-{end} after {max_retries} attempts (URL kept expiring)"
+        )
+    
+    def _load_full(self):
+        """Load the entire attachment into memory buffer using read_attachment()."""
+        if self._fully_loaded:
+            return
+        
+        from octostar.utils.workspace import read_attachment
+        
+        data = read_attachment.sync(
+            os_workspace=self._workspace_id,
+            os_entity_uid=self._entity_id,
+            decode=False,
+            stream=False,
+            client=self._client
+        )
+        self._buffer = BytesIO(data or b"")
+        self._fully_loaded = True
+        self._size = len(data) if data else 0
+        self._position = 0
+    
+    def read(self, size: int = -1) -> bytes:
+        if self._fully_loaded:
+            if not self._buffer:
+                self._load_full()
+            return self._buffer.read(size)
+        
+        # Otherwise, use HTTP Range requests for efficient streaming
+        if size == -1:
+            # Read all from current position
+            self._fetch_size()
+            if self._position >= self._size:
+                return b""
+            size = self._size - self._position
+        
+        if size <= 0:
+            return b""
+        
+        # Fetch the size if we don't know it
+        self._fetch_size()
+        
+        # Adjust size if it goes beyond the end
+        if self._position + size > self._size:
+            size = self._size - self._position
+        
+        if size <= 0:
+            return b""
+        
+        # Read using HTTP Range request
+        end_byte = self._position + size - 1
+        data = self._read_range(self._position, end_byte)
+        self._position += len(data)
+        
+        return data
+    
+    def write(self, b: bytes) -> int:
+        if not self._buffer:
+            self._load_full()
+        
+        n = self._buffer.write(b)
+        self._modified = True
+        return n
+    
+    def seek(self, offset: int, whence: int = SEEK_SET) -> int:
+        # Calculate new position
+        if whence == SEEK_SET:
+            new_pos = offset
+        elif whence == SEEK_CUR:
+            if self._buffer and self._fully_loaded:
+                new_pos = self._buffer.tell() + offset
+            else:
+                new_pos = self._position + offset
+        elif whence == SEEK_END:
+            self._fetch_size()
+            new_pos = self._size + offset
+        else:
+            raise ValueError(f"Invalid whence value: {whence}")
+        
+        if new_pos < 0:
+            raise ValueError("Negative seek position")
+        
+        # If fully loaded, use buffer seek
+        if self._buffer and self._fully_loaded:
+            return self._buffer.seek(new_pos, SEEK_SET)
+        
+        # Otherwise, just update position (HTTP Range will handle it)
+        self._position = new_pos
+        return self._position
+    
+    def tell(self) -> int:
+        if self._buffer and self._fully_loaded:
+            return self._buffer.tell()
+        return self._position
+    
+    def flush(self):
+        """Flush the internal buffer and write to workspace if modified."""
+        if self._buffer:
+            self._buffer.flush()
+        
+        if not self._modified or not self._buffer:
+            return
+        
+        if not self._entity_type or not self._filetype:
+            raise ValueError("entity_type and filetype required to flush to workspace")
+        
+        from octostar.utils.workspace import write_attachment
+        
+        current_pos = self._buffer.tell()
+        self._buffer.seek(0, SEEK_SET)
+        data = self._buffer.read()
+        self._buffer.seek(current_pos, SEEK_SET)
+        
+        write_attachment.sync(
+            os_workspace=self._workspace_id,
+            os_entity_uid=self._entity_id,
+            entity_type=self._entity_type,
+            filetype=self._filetype,
+            file=data,
+            client=self._client
+        )
+        self._modified = False
+    
+    def close(self):
+        if not self._closed:
+            if self._modified:
+                self.flush()
+            if self._buffer:
+                self._buffer.close()
+            if self._http_client:
+                self._http_client.close()
+                self._http_client = None
+            super().close()
+    
+    def delete(self):
+        """Delete the entity from the workspace using delete_entity()."""
+        from octostar.utils.workspace import delete_entity
+        
+        delete_entity.sync(
+            os_entity_uid=self._entity_id,
+            client=self._client
+        )
+    
+    def truncate(self, size: Optional[int] = None) -> int:
+        if not self._buffer:
+            self._load_full()
+        self._modified = True
+        return self._buffer.truncate(size)
+    
+    def getvalue(self) -> bytes:
+        if not self._buffer or not self._fully_loaded:
+            self._load_full()
+        return self._buffer.getvalue()
+    
+    def to_locator(self) -> Dict[str, Any]:
+        """
+        Serialize to locator with workspace and entity ID.
+        
+        Returns:
+            {"location": "workspace_attachment", "pointer": "workspace_id/entity_id", 
+             "entity_type": "...", "filetype": "..."}
+        """
+        if self._workspace_id and self._entity_id:
+            pointer = f"{self._workspace_id}/{self._entity_id}"
+        else:
+            pointer = None
+        
+        locator = {
+            "location": ContentsLocation.WORKSPACE_ATTACHMENT.value,
+            "pointer": pointer
+        }
+        if self._entity_type:
+            locator["entity_type"] = self._entity_type
+        if self._filetype:
+            locator["filetype"] = self._filetype
+        return locator
+    
+    @staticmethod
+    def _from_locator(locator: Dict[str, Any], client=None) -> 'WorkspaceAttachmentContents':
+        """
+        Create WorkspaceAttachmentContents from a locator dictionary.
+        
+        Args:
+            locator: Locator dictionary with pointer "workspace_id/entity_id"
+            client: Octostar client for remote operations
+            
+        Returns:
+            New WorkspaceAttachmentContents instance
+        """
+        pointer = locator.get("pointer")
+        workspace_id = None
+        entity_id = None
+        
+        if pointer:
+            parts = pointer.split("/")
+            if len(parts) >= 2:
+                workspace_id = parts[0]
+                entity_id = parts[-1]
+        
+        return WorkspaceAttachmentContents(
+            entity_type=locator.get("entity_type"),
+            filetype=locator.get("filetype"),
+            workspace_id=workspace_id,
+            entity_id=entity_id,
+            client=client
+        )
+
+
+class TemporaryAttachmentContents(Contents):
+    """
+    Contents implementation for Octostar temporary blob storage.
+    
+    Uses octostar-api utilities (read_temporary_blob, write_temporary_blob,
+    delete_temporary_blob) to store files in the user's temporary S3 bucket.
+    
+    Temporary blobs are keyed by filename (not workspace/entity), and are not
+    associated with any workspace entity. Use WorkspaceAttachmentContents for that.
+    """
+    
+    def __init__(
+        self,
+        entity_type: Optional[str] = None,
+        filetype: Optional[str] = None,
+        *,
+        filename: str,
+        client,
+        initial_data: Optional[bytes] = None,
+        **kwargs
+    ):
+        super().__init__(entity_type, filetype, **kwargs)
+        self._filename = filename
+        self._client = client
+        
+        self._buffer: Optional[BytesIO] = None
+        self._fully_loaded = False
+        self._modified = False
+        
+        if initial_data is not None:
+            self._buffer = BytesIO(initial_data)
+            self._fully_loaded = True
+    
+    def _load_full(self):
+        """Load the entire blob into memory using read_temporary_blob()."""
+        if self._fully_loaded:
+            return
+        
+        from octostar.utils.workspace import read_temporary_blob
+        
+        data = read_temporary_blob.sync(
+            filename=self._filename,
+            decode=False,
+            client=self._client
+        )
+        self._buffer = BytesIO(data or b"")
+        self._fully_loaded = True
+    
+    def read(self, size: int = -1) -> bytes:
+        if not self._buffer:
+            self._load_full()
+        return self._buffer.read(size)
+    
+    def write(self, b: bytes) -> int:
+        if not self._buffer:
+            self._load_full()
+        n = self._buffer.write(b)
+        self._modified = True
+        return n
+    
+    def seek(self, offset: int, whence: int = SEEK_SET) -> int:
+        if not self._buffer:
+            self._load_full()
+        return self._buffer.seek(offset, whence)
+    
+    def tell(self) -> int:
+        if not self._buffer:
+            self._load_full()
+        return self._buffer.tell()
+    
+    def flush(self):
+        """Flush the internal buffer and write to temp bucket if modified."""
+        if self._buffer:
+            self._buffer.flush()
+        
+        if not self._modified or not self._buffer:
+            return
+        
+        from octostar.utils.workspace import write_temporary_blob
+        
+        current_pos = self._buffer.tell()
+        self._buffer.seek(0, SEEK_SET)
+        data = self._buffer.read()
+        self._buffer.seek(current_pos, SEEK_SET)
+        
+        write_temporary_blob.sync(
+            filename=self._filename,
+            file=data,
+            client=self._client
+        )
+        self._modified = False
+    
+    def close(self):
+        if not self._closed:
+            if self._modified:
+                self.flush()
+            if self._buffer:
+                self._buffer.close()
+            super().close()
+    
+    def delete(self):
+        """Delete the blob from the temporary bucket."""
+        from octostar.utils.workspace import delete_temporary_blob
+        
+        delete_temporary_blob.sync(
+            filename=self._filename,
+            client=self._client
+        )
+    
+    def truncate(self, size: Optional[int] = None) -> int:
+        if not self._buffer:
+            self._load_full()
+        self._modified = True
+        return self._buffer.truncate(size)
+    
+    def getvalue(self) -> bytes:
+        if not self._buffer or not self._fully_loaded:
+            self._load_full()
+        return self._buffer.getvalue()
+    
+    def to_locator(self) -> Dict[str, Any]:
+        """
+        Serialize to locator with filename.
+        
+        Returns:
+            {"location": "temporary_attachment", "filename": "..."}
+        """
+        locator = {
+            "location": ContentsLocation.TEMPORARY_ATTACHMENT.value,
+            "filename": self._filename
+        }
+        if self._entity_type:
+            locator["entity_type"] = self._entity_type
+        if self._filetype:
+            locator["filetype"] = self._filetype
+        return locator
+    
+    @staticmethod
+    def _from_locator(locator: Dict[str, Any], client=None) -> 'TemporaryAttachmentContents':
+        """
+        Create TemporaryAttachmentContents from a locator dictionary.
+        
+        Args:
+            locator: Locator dictionary with filename
+            client: Octostar client
+            
+        Returns:
+            New TemporaryAttachmentContents instance
+        """
+        return TemporaryAttachmentContents(
+            entity_type=locator.get("entity_type"),
+            filetype=locator.get("filetype"),
+            filename=locator["filename"],
+            client=client
+        )

{streamlit_octostar_utils-0.4.1 → streamlit_octostar_utils-0.4.2.dev2}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -26,6 +26,7 @@ from ..core.dict import recursive_update_dict, travel_dict, jsondict_hash
 from ..core.timestamp import now, string_to_datetime
 from .fastapi import DefaultErrorRoute, Route
 from ..ontology.inheritance import is_child_concept as is_child_concept_fn, get_label_keys
+from .contents import Contents, MemoryContents, WorkspaceAttachmentContents, TemporaryAttachmentContents, ContentsLocation
 
 RELATIONSHIP_ENTITY_NAME = "os_relationship"
 LOCAL_RELATIONSHIP_ENTITY_NAME = "os_workspace_relationship"
@@ -47,11 +48,6 @@ def safe_async_run(coro):
         return asyncio.run(coro)
 
 
-class NifiContentsPointerLocationModel(Enum):
-    LOCAL = "local"
-    ATTACHMENT = "attachment"
-
-
 class NifiProxyEntityModel(BaseModel):
     entity_id: str
     entity_type: str
@@ -72,10 +68,6 @@ class NifiEntityModel(BaseModel):
             relationships: List[str]
             label_keys: List[str]
 
-        class ContentsPointerModel(BaseModel):
-            location: NifiContentsPointerLocationModel
-            pointer: Optional[str] = None
-
         jwt: str
         ontology_name: str
         ontology_info: OntologyInfoModel
@@ -84,7 +76,7 @@ class NifiEntityModel(BaseModel):
         nifi_attributes: dict = Field(default_factory=dict)
         config: dict = Field(default_factory=dict)
         metrics: dict = Field(default_factory=dict)
-        contents_pointer: Optional[ContentsPointerModel] = None
+        contents_pointer: Optional[dict] = None
         is_temporary: bool = False
         exception: dict = Field(default_factory=dict)
         last_processor_name: Optional[str] = None
@@ -103,7 +95,7 @@ class NifiEntityModel(BaseModel):
     record: RecordModel
     annotations: Dict[str, Any] = Field(default_factory=dict)
     children: List[Union[NifiOTMRelationshipProxyModel, NifiProxyEntityModel]] = []
-    contents: Optional[bytes] = None
+    contents: Optional[Dict[str, Any]] = None
 
 
 NifiEntityModel.model_rebuild()
@@ -295,12 +287,10 @@ class NifiContextManager(object):
         return client, curr_user_ontology
 
     def receive_input(self, json_data, processor_name) -> List["NifiEntityBatch"]:
-        def _safe_decode(contents):
-            return base64.b64decode(contents) if contents else None
-
         entities = []
         all_independent_uids = [e["record"]["entity_id"] for e in json_data]
         for elem in json_data:
+            contents = Contents.from_locator(elem.get("contents"), client=self.client)
             entities.append(
                 NifiEntity(
                     self,
@@ -309,7 +299,7 @@ class NifiContextManager(object):
                     elem["annotations"],
                     all_independent_uids,
                     elem["children"],
-                    _safe_decode(elem.get("contents")),
+                    contents,
                 )
             )
         entities = sorted(
@@ -463,10 +453,13 @@ class NifiContextManager(object):
         fetch_concept_relationships = {}
         # FIND FILES TO WRITE
         for entity in entities:
-            if entity.is_child_concept("os_file"):
+            if entity.is_child_concept("os_attachable"):
                 has_write_flag = entity.sync_params.get(NifiContextManager.SyncFlag.WRITE_CONTENTS)
-                is_temp_with_pointer = entity.request.get("is_temporary") and entity.contents_pointer
-                if has_write_flag or is_temp_with_pointer:
+                is_temporary = entity.request.get("is_temporary")
+                if is_temporary:
+                    if entity._contents and not isinstance(entity._contents, WorkspaceAttachmentContents):
+                        files_to_write.append(entity)
+                elif has_write_flag:
                     if entity.contents:
                         files_to_write.append(entity)
         # FIND ENTITIES TO UPSERT
@@ -495,6 +488,9 @@ class NifiContextManager(object):
             # WRITE FILES
             if files_to_write:
                 for file in files_to_write:
+                    if not file.contents:
+                        continue
+                    old_contents = file._contents
                     new_file_record = write_file.sync(
                         file.write_os_workspace,
                         "./" + file.record["os_item_name"],
@@ -510,11 +506,24 @@ class NifiContextManager(object):
                     file.record["entity_label"] = file.label
                     file.request["is_temporary"] = False
                     file.request["entity_timestamp"] = file.record["os_last_updated_at"]
-                    file._contents = None
+                    file._contents = WorkspaceAttachmentContents(
+                        workspace_id=file.record['os_workspace'],
+                        entity_id=file.record['os_entity_uid'],
+                        client=self.client,
+                        entity_type=file.record["os_concept"],
+                        filetype=file.record["os_item_content_type"]
+                    )
                     file.request["contents_pointer"] = {
-                        "location": NifiContentsPointerLocationModel.ATTACHMENT.value,
-                        "pointer": f"{file.record['os_workspace']}/{file.record['os_entity_uid']}"
+                        "location": ContentsLocation.WORKSPACE_ATTACHMENT.value,
+                        "pointer": f"{file.record['os_workspace']}/{file.record['os_entity_uid']}",
+                        "entity_type": file.record["os_concept"],
+                        "filetype": file.record["os_item_content_type"]
                     }
+                    if isinstance(old_contents, TemporaryAttachmentContents):
+                        try:
+                            old_contents.delete()
+                        except Exception:
+                            pass
             # UPSERT ENTITIES
             if entities_to_upsert:
                 new_entities = upsert_entities.sync(
@@ -692,7 +701,7 @@ class NifiEntity(object):
                 c["annotations"],
                 all_independent_uids,
                 c["children"],
-                c["contents"],
+                Contents.from_locator(c.get("contents"), client=self.context.client),
             )
             for c in full_entity_children
         ]
@@ -711,7 +720,7 @@ class NifiEntity(object):
             for i in range(len(child_uids))
         ]
         self.children.extend(proxy_otm_children)
-        self._contents = contents
+        self._contents: Optional[Contents] = contents
         self.drop_on_output = False
 
     def __eq__(self, other):
@@ -741,36 +750,16 @@ class NifiEntity(object):
         self._annotations = new_annotations
 
     @property
-    def contents(self):
+    def contents(self) -> Optional[Contents]:
         if not self._contents:
             contents_pointer = self.contents_pointer
             if not contents_pointer:
                 return None
-            if contents_pointer["location"] == "attachment":
-                self._contents = read_file.sync(
-                    contents_pointer["pointer"].split("/")[0],
-                    contents_pointer["pointer"].split("/")[-1],
-                    False,
-                    client=self.context.client,
-                )
+            self._contents = Contents.from_locator(contents_pointer, client=self.context.client)
         return self._contents
 
-    @property
-    def contents_pointer(self):
-        contents_pointer = deepcopy(self.request.get("contents_pointer"))
-        if not self.request.get("contents_pointer"):
-            return None
-        ptr_location = contents_pointer.get("location")
-        if ptr_location == "attachment" and not contents_pointer.get("pointer"):
-            contents_pointer["pointer"] = f"{self.record['os_workspace']}/{self.record['os_entity_uid']}"
-        return contents_pointer
-
-    @contents_pointer.setter
-    def contents_pointer(self, new_value):
-        self.request["contents_pointer"] = new_value
-
     @contents.setter
-    def contents(self, new_contents):
+    def contents(self, new_contents: Optional[Union[Contents, bytes]]):
         self._contents = new_contents
 
     @property
@@ -866,9 +855,6 @@ class NifiEntity(object):
         return not _is_sub_fragment_recursive(fragment)
 
     def to_json(self):
-        def _safe_encode(contents):
-            return base64.b64encode(contents) if contents else None
-
         if self.drop_on_output:
             return
         proxy_entity_children = []
@@ -909,7 +895,7 @@ class NifiEntity(object):
             "record": self.record,
             "children": children,
             "annotations": self.annotations,
-            "contents": _safe_encode(self._contents),
+            "contents": self._contents.to_locator() if self._contents else None,
         }
 
     def _add_entity(self, os_workspace, entity_type, fields, os_entity_uid=None):
@@ -1058,7 +1044,7 @@ class NifiEntity(object):
         os_parent_folder,
         filename,
         filetype,
-        file,
+        file: Union[Contents, bytes],
         fields={},
         os_relationship_name=FILE_RELATIONSHIP_NAME,
         os_relationship_type="mtm",
@@ -1080,8 +1066,20 @@ class NifiEntity(object):
             os_entity_uid,
             os_relationship_uid,
         )
-        child_entity._contents = file
-        child_entity.request["contents_pointer"] = NifiEntityModel.RequestModel.ContentsPointerModel(location="local")
+        if isinstance(file, Contents):
+            child_entity._contents = file
+        else:
+            temp_filename = f"tmp_{child_entity.record['os_entity_uid']}"
+            temp_contents = TemporaryAttachmentContents(
+                entity_type=FILE_ENTITY_NAME,
+                filetype=filetype,
+                filename=temp_filename,
+                client=self.context.client,
+                initial_data=file,
+            )
+            temp_contents.flush()
+            child_entity._contents = temp_contents
+        child_entity.request["contents_pointer"] = child_entity._contents.to_locator()
         return child_entity, child_rel
 
     def add_tag(self, os_workspace, name, group, order, color, fields={}):

{streamlit_octostar_utils-0.4.1 → streamlit_octostar_utils-0.4.2.dev2}/streamlit_octostar_utils/core/timestamp.py

@@ -10,7 +10,10 @@ def now():
 
 
 def string_to_datetime(datetime_str):
-    return dt_parser.parse(
+    parsed = dt_parser.parse(
         datetime_str
         or dt.datetime.fromtimestamp(0, dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
     )
+    if parsed.tzinfo is None or parsed.tzinfo.utcoffset(parsed) is None:
+        parsed = parsed.replace(tzinfo=dt.timezone.utc)
+    return parsed