cognite-toolkit 0.5.44__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/utils/modules.py

@@ -4,7 +4,6 @@ from pathlib import Path
 from typing import Literal, overload
 
 from cognite_toolkit._cdf_tk.constants import (
-    BUILTIN_MODULES,
     MODULE_PATH_SEP,
     ROOT_MODULES,
 )
@@ -23,9 +22,6 @@ def iterate_modules(root_dir: Path) -> Iterator[tuple[Path, list[Path]]]:
     if root_dir.name in ROOT_MODULES:
         yield from _iterate_modules(root_dir)
         return
-    elif root_dir.name == BUILTIN_MODULES:
-        yield from _iterate_modules(root_dir)
-        return
     for root_module in ROOT_MODULES:
         module_dir = root_dir / root_module
         if module_dir.exists():
@@ -35,7 +31,7 @@ def iterate_modules(root_dir: Path) -> Iterator[tuple[Path, list[Path]]]:
 def _iterate_modules(root_dir: Path) -> Iterator[tuple[Path, list[Path]]]:
     # local import to avoid circular import
     from cognite_toolkit._cdf_tk.constants import EXCL_FILES
-    from cognite_toolkit._cdf_tk.loaders import LOADER_BY_FOLDER_NAME
+    from cognite_toolkit._cdf_tk.cruds import CRUDS_BY_FOLDER_NAME
 
     if not root_dir.exists():
         return
@@ -43,7 +39,7 @@ def _iterate_modules(root_dir: Path) -> Iterator[tuple[Path, list[Path]]]:
         if not module_dir.is_dir():
             continue
         sub_directories = [path for path in module_dir.iterdir() if path.is_dir()]
-        is_any_resource_directories = any(dir.name in LOADER_BY_FOLDER_NAME for dir in sub_directories)
+        is_any_resource_directories = any(dir.name in CRUDS_BY_FOLDER_NAME for dir in sub_directories)
         if sub_directories and is_any_resource_directories:
             # Module found
             yield module_dir, [path for path in module_dir.rglob("*") if path.is_file() and path.name not in EXCL_FILES]
@@ -63,13 +59,13 @@ def module_from_path(path: Path, return_resource_folder: Literal[False] = False)
 def module_from_path(path: Path, return_resource_folder: bool = False) -> str | tuple[str, str]:
     """Get the module name from a path"""
     # local import to avoid circular import
-    from cognite_toolkit._cdf_tk.loaders import LOADER_BY_FOLDER_NAME
+    from cognite_toolkit._cdf_tk.cruds import CRUDS_BY_FOLDER_NAME
 
     if len(path.parts) == 1:
         raise ValueError("Path is not a module")
     last_folder = path.parts[1]
     for part in path.parts[1:]:
-        if part in LOADER_BY_FOLDER_NAME:
+        if part in CRUDS_BY_FOLDER_NAME:
             if return_resource_folder:
                 return last_folder, part
             return last_folder
@@ -80,13 +76,13 @@ def module_from_path(path: Path, return_resource_folder: bool = False) -> str |
 def module_directory_from_path(path: Path) -> Path:
     """Get the module directory from a path"""
     # local import to avoid circular import
-    from cognite_toolkit._cdf_tk.loaders import LOADER_BY_FOLDER_NAME
+    from cognite_toolkit._cdf_tk.cruds import CRUDS_BY_FOLDER_NAME
 
     if len(path.parts) == 1:
         raise ValueError("Path is not a module")
 
     for _ in range(len(path.parts)):
-        if path.name in LOADER_BY_FOLDER_NAME and path.parent != path:
+        if path.name in CRUDS_BY_FOLDER_NAME and path.parent != path:
             return path.parent
         path = path.parent
     raise ValueError("Path is not part of a module")
@@ -95,10 +91,10 @@ def module_directory_from_path(path: Path) -> Path:
 def resource_folder_from_path(path: Path) -> str:
     """Get the resource_folder from a path"""
     # local import to avoid circular import
-    from cognite_toolkit._cdf_tk.loaders import LOADER_BY_FOLDER_NAME
+    from cognite_toolkit._cdf_tk.cruds import CRUDS_BY_FOLDER_NAME
 
     for part in path.parts:
-        if part in LOADER_BY_FOLDER_NAME:
+        if part in CRUDS_BY_FOLDER_NAME:
             return part
     raise ValueError("Path does not contain a resource folder")
 
@@ -106,12 +102,12 @@ def resource_folder_from_path(path: Path) -> str:
 def is_module_path(path: Path) -> bool:
     """Check if a path is a module path"""
     # local import to avoid circular import
-    from cognite_toolkit._cdf_tk.loaders import LOADER_BY_FOLDER_NAME
+    from cognite_toolkit._cdf_tk.cruds import CRUDS_BY_FOLDER_NAME
 
     if not path.is_dir():
         return False
 
-    return any(sub_folder.name in LOADER_BY_FOLDER_NAME for sub_folder in path.iterdir() if sub_folder.is_dir())
+    return any(sub_folder.name in CRUDS_BY_FOLDER_NAME for sub_folder in path.iterdir() if sub_folder.is_dir())
 
 
 def find_directory_with_subdirectories(

cognite_toolkit/_cdf_tk/utils/producer_worker.py

@@ -1,123 +1,340 @@
 import queue
+import sys
 import threading
+import time
+import traceback
+import typing
 from collections.abc import Callable, Iterable, Sized
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 from rich.console import Console
-from rich.progress import BarColumn, Progress, TaskID, TaskProgressColumn, TextColumn, TimeRemainingColumn
+from rich.panel import Panel
+from rich.progress import (
+    BarColumn,
+    Progress,
+    ProgressColumn,
+    SpinnerColumn,
+    Task,
+    TaskID,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+
+from cognite_toolkit._cdf_tk.exceptions import ToolkitRuntimeError
 
 T_Download = TypeVar("T_Download", bound=Sized)
 T_Processed = TypeVar("T_Processed", bound=Sized)
+T_Item = TypeVar("T_Item", bound=Sized)
+
+# Sentinels for signaling finish
+PROCESS_FINISH_SENTINEL = object()
+WRITE_FINISH_SENTINEL = object()
+
+
+class ItemCountColumn(ProgressColumn):
+    def render(self, task: Task) -> str:
+        return f"[green]{int(task.fields.get('item_count', 0)):,} items"
 
 
 class ProducerWorkerExecutor(Generic[T_Download, T_Processed]):
+    """This class manages the execution of a producer-worker pattern with multiple threads.
+
+    It downloads data from an iterable, processes it, and writes the processed data.
+
+    Args:
+        download_iterable (Iterable[T_Download]): An iterable that yields chunks of items.
+        process (Callable[[T_Download], T_Processed]): A function that processes the chunks of items.
+        write (Callable[[T_Processed], None]): A function that writes the processed items. This is typically to
+            a file, CDF, or another storage format.
+        iteration_count (int): The total number of iterations expected. Note this is not the total number of items,
+            but rather the number of chunks of items that will be processed. Set to `None` if the total number of
+            iterations is unknown.
+        max_queue_size (int): The maximum size of the queues used for processing and writing. Note that the queue
+            holds chunks of items, not individual items, so this should be set based on the expected size of the
+            chunks being processed.
+        download_description (str): A description of the download process, used for progress tracking.
+        process_description (str): A description of the processing step, used for progress tracking.
+        write_description (str): A description of the writing step, used for progress tracking.
+        console (Console | None): An optional Rich Console instance for outputting progress and error messages.
+
+    Examples:
+        >>> from cognite_toolkit._cdf_tk.utils.producer_worker import ProducerWorkerExecutor
+        >>> from typing import List
+        >>> def download_data() -> Iterable[list[int]]:
+        ...     yield [1, 2, 3, 4, 5]
+        ...     yield [6, 7, 8, 9, 10]
+        >>> def process_data(data: list[int]) -> list[int]:
+        ...     return [x * 2 for x in data]
+        >>> def write_data(data: list[int]) -> None:
+        ...     print(f"Writing data: {data}")
+        >>> executor = ProducerWorkerExecutor(
+        ...     download_iterable=download_data(),
+        ...     process=process_data,
+        ...     write=write_data,
+        ...     iteration_count=2,
+        ...     max_queue_size=1,
+        ... )
+        >>> executor.run()
+
+    """
+
     def __init__(
         self,
         download_iterable: Iterable[T_Download],
         process: Callable[[T_Download], T_Processed],
-        write_to_file: Callable[[T_Processed], None],
-        iteration_count: int,
+        write: Callable[[T_Processed], None],
+        iteration_count: int | None,
         max_queue_size: int,
+        download_description: str = "Producing",
+        process_description: str = "Processing",
+        write_description: str = "Writing",
+        console: Console | None = None,
+        verbose: bool = False,
     ) -> None:
         self._download_iterable = download_iterable
-        self.download_complete = False
-        self.is_processing = False
-        self._write_to_file = write_to_file
         self._process = process
-        self.console = Console()
-        self.process_queue: queue.Queue[T_Download] = queue.Queue(maxsize=max_queue_size)
-        self.file_queue: queue.Queue[T_Processed] = queue.Queue(maxsize=max_queue_size)
-
+        self._write = write
         self.iteration_count = iteration_count
+        self.download_description = download_description
+        self.process_description = process_description
+        self.write_description = write_description
+        self.console = console or Console()
+        self._stop_event = threading.Event()
+        self._error_event = threading.Event()
+        # Queues for managing the flow of data between threads
+        # Download -> [process_queue] -> Process -> [write_queue] -> Write
+        self.process_queue: queue.Queue[T_Download] = queue.Queue(maxsize=max_queue_size)
+        self.write_queue: queue.Queue[T_Processed] = queue.Queue(maxsize=max_queue_size)
         self.total_items = 0
-        self.error_occurred = False
         self.error_message = ""
+        self.error_traceback = ""
+        self.verbose = verbose
+
+    @property
+    def error_occurred(self) -> bool:
+        """Indicates if an error occurred during execution."""
+        return self._error_event.is_set()
+
+    @property
+    def stopped_by_user(self) -> bool:
+        """Indicates if the execution was stopped by the user."""
+        return self._stop_event.is_set()
+
+    def _get_progress_columns(self) -> list[ProgressColumn]:
+        """Helper to set up the progress bar and tasks based on iteration_count."""
+        if self.iteration_count is None:
+            return [
+                SpinnerColumn(),
+                TextColumn("[bold blue]{task.description}"),
+                ItemCountColumn(),
+                TimeElapsedColumn(),
+            ]
+        else:
+            return [
+                TextColumn("[bold blue]{task.description}"),
+                BarColumn(),
+                TaskProgressColumn(),
+                TimeRemainingColumn(),
+            ]
+
+    def _user_input_listener(self, producer_thread: threading.Thread) -> None:
+        while not self._error_event.is_set() and not self._stop_event.is_set():
+            key = getch(timeout=0.1)
+            if key is None and not producer_thread.is_alive():
+                break
+            elif key is None:
+                continue
+            elif key.casefold() == "q":
+                self._stop_event.set()
+                self.console.print(
+                    f"[yellow]Execution stopped by user. Finishing: {self.write_description}...[/yellow]"
+                )
+                break
 
     def run(self) -> None:
-        with Progress(
-            TextColumn("[bold blue]{task.description}"),
-            BarColumn(),
-            TaskProgressColumn(),
-            TimeRemainingColumn(),
-            console=self.console,
-        ) as progress:
-            download_task = progress.add_task("Downloading", total=self.iteration_count)
+        self.console.print(f"[blue]Starting {self.download_description} (Press 'q' to stop)...[/blue]")
+        columns = self._get_progress_columns()
+        with Progress(*columns, console=self.console) as progress:
+            task_args: dict[str, Any] = (
+                {"item_count": 0, "total": None} if self.iteration_count is None else {"total": self.iteration_count}
+            )
+            download_task = progress.add_task(self.download_description, **task_args)
+            process_task = progress.add_task(self.process_description, **task_args)
+            write_task = progress.add_task(self.write_description, **task_args)
+
             download_thread = threading.Thread(target=self._download_worker, args=(progress, download_task))
-            process_task = progress.add_task("Processing", total=self.iteration_count)
             process_thread = threading.Thread(target=self._process_worker, args=(progress, process_task))
-
-            write_task = progress.add_task("Writing to file", total=self.iteration_count)
             write_thread = threading.Thread(target=self._write_worker, args=(progress, write_task))
 
             download_thread.start()
             process_thread.start()
             write_thread.start()
 
-            # Wait for all threads to finish
-            download_thread.join()
-            process_thread.join()
-            write_thread.join()
+            input_thread = threading.Thread(target=self._user_input_listener, args=(download_thread,))
+            input_thread.start()
+
+            for t in [download_thread, process_thread, write_thread]:
+                try:
+                    t.join()
+                except KeyboardInterrupt:
+                    self.console.print("[red]Execution interrupted by user.[/red]")
+                    self._stop_event.set()
+                    break
+
+            # After a possible interrupt, we must wait for all threads to finish their
+            # graceful shutdown. This is important to prevent data loss.
+            for t in [download_thread, process_thread, write_thread, input_thread]:
+                if t.is_alive():
+                    t.join()
+
+    def raise_on_error(self) -> None:
+        """Raises an exception if an error occurred during execution."""
+        if self._error_event.is_set():
+            if self.verbose and self.error_traceback:
+                self.print_traceback()
+            raise ToolkitRuntimeError(f"An error occurred during execution: {self.error_message}")
+        if self._stop_event.is_set():
+            raise ToolkitRuntimeError("Execution was stopped by the user.")
+
+    def print_traceback(self) -> None:
+        """Prints the traceback if an error occurred during execution."""
+        if self.error_traceback:
+            self.console.print(
+                Panel(
+                    self.error_traceback,
+                    title="Traceback",
+                    expand=False,
+                    border_style="red",
+                )
+            )
 
     def _download_worker(self, progress: Progress, download_task: TaskID) -> None:
         """Worker thread for downloading data."""
-        iterator = iter(self._download_iterable)
-        while not self.error_occurred:
+        try:
+            iterator = iter(self._download_iterable)
+        except TypeError as e:
+            self._error_event.set()
+            self.error_message = str(e)
+            self.console.print(f"[red]Error[/red] occurred while {self.download_description}: {self.error_message}")
+            return
+        item_count = 0
+        while not self._error_event.is_set():
             try:
+                if self._stop_event.is_set():
+                    break
                 items = next(iterator)
                 self.total_items += len(items)
-                while not self.error_occurred:
-                    try:
-                        self.process_queue.put(items, timeout=0.5)
-                        progress.update(download_task, advance=1)
-                        break  # Exit the loop once the item is successfully added
-                    except queue.Full:
-                        # Retry until the queue has space
-                        continue
+                if self._put_with_error_check(items, self.process_queue):
+                    item_count += len(items)
+                    progress.update(download_task, advance=1, item_count=item_count)
+                    continue
+                break  # Exit if error event was set while waiting to put
             except StopIteration:
-                self.download_complete = True
                 break
             except Exception as e:
-                self.error_occurred = True
-                self.error_message = str(e)
-                self.console.print(f"[red]Error[/red] occurred while downloading: {self.error_message}")
+                self._error_event.set()
+                self.error_message = f"{type(e).__name__} {e!s}"
+                self.error_traceback = traceback.format_exc()
+                self.console.print(f"[red]Error[/red] occurred while {self.download_description}: {self.error_message}")
                 break
+        self._put_with_error_check(PROCESS_FINISH_SENTINEL, self.process_queue)  # type: ignore[misc]
+
+    def _put_with_error_check(self, items: T_Item, target_queue: queue.Queue[T_Item]) -> bool:
+        """Helper to put items into a queue with error checking."""
+        while not self._error_event.is_set():
+            try:
+                target_queue.put(items, timeout=0.5)
+                return True
+            except queue.Full:
+                continue
+        return False
 
     def _process_worker(self, progress: Progress, process_task: TaskID) -> None:
         """Worker thread for processing data."""
-        self.is_processing = True
-        while (not self.download_complete or not self.process_queue.empty()) and not self.error_occurred:
+        item_count = 0
+        while not self._error_event.is_set():
             try:
                 items = self.process_queue.get(timeout=0.5)
+                if items is PROCESS_FINISH_SENTINEL:
+                    # Signal writer to finish
+                    self._put_with_error_check(WRITE_FINISH_SENTINEL, self.write_queue)  # type: ignore[misc]
+                    self.process_queue.task_done()
+                    break
                 processed_items = self._process(items)
-                self.file_queue.put(processed_items)
-                progress.update(process_task, advance=1)
-                self.process_queue.task_done()
+                if self._put_with_error_check(processed_items, self.write_queue):
+                    item_count += len(items)
+                    progress.update(process_task, advance=1, item_count=item_count)
+                    self.process_queue.task_done()
+                    continue
+                else:
+                    self.process_queue.task_done()
+                    break  # Exit if error event was set while waiting to put
             except queue.Empty:
                 continue
             except Exception as e:
-                self.error_occurred = True
-                self.error_message = str(e)
-                self.console.print(f"[red]ErrorError[/red] occurred while processing: {self.error_message}")
+                self._error_event.set()
+                self.error_message = f"{type(e).__name__} {e!s}"
+                self.error_traceback = traceback.format_exc()
+                self.console.print(f"[red]Error[/red] occurred while {self.process_description}: {self.error_message}")
                 break
-        self.is_processing = False
 
     def _write_worker(self, progress: Progress, write_task: TaskID) -> None:
         """Worker thread for writing data to file."""
-        while (
-            not self.download_complete
-            or self.is_processing
-            or not self.file_queue.empty()
-            or not self.process_queue.empty()
-        ) and not self.error_occurred:
+        item_count = 0
+        while not self._error_event.is_set():
             try:
-                items = self.file_queue.get(timeout=0.5)
-                self._write_to_file(items)
-                progress.update(write_task, advance=1)
-                self.file_queue.task_done()
+                items = self.write_queue.get(timeout=0.5)
+                if items is WRITE_FINISH_SENTINEL:
+                    self.write_queue.task_done()
+                    break
+                self._write(items)
+                item_count += len(items)
+                progress.update(write_task, advance=1, item_count=item_count)
+                self.write_queue.task_done()
             except queue.Empty:
                 continue
             except Exception as e:
-                self.error_occurred = True
-                self.error_message = str(e)
-                self.console.print(f"[red]Error[/red] occurred while writing: {self.error_message}")
+                self._error_event.set()
+                self.error_message = f"{type(e).__name__} {e!s}"
+                self.error_traceback = traceback.format_exc()
+                self.console.print(f"[red]Error[/red] occurred while {self.write_description}: {self.error_message}")
                 break
+
+
+# MyPy fails as the imports are os specific
+# thus we disable type checking for this function
+@typing.no_type_check
+def getch(timeout: float) -> str | None:
+    """Get a single character from standard input. Does not echo to the screen."""
+    if not sys.stdin.isatty():
+        return None
+    try:
+        # Windows
+        import msvcrt
+
+        end_time = time.time() + timeout
+        while time.time() < end_time:
+            if msvcrt.kbhit():
+                return msvcrt.getwch()
+            time.sleep(0.01)
+        return None
+    except ImportError:
+        # Unix
+        import select
+        import termios
+        import tty
+
+        fd = sys.stdin.fileno()
+        old_settings = termios.tcgetattr(fd)
+        try:
+            tty.setraw(fd)
+            rlist, _, _ = select.select([fd], [], [], timeout)
+            if rlist:
+                ch = sys.stdin.read(1)
+                return ch
+            else:
+                return None
+        finally:
+            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)

cognite_toolkit/_cdf_tk/utils/progress_tracker.py

@@ -0,0 +1,87 @@
+import threading
+from collections import Counter
+from typing import Generic, Literal, TypeAlias, get_args, overload
+
+from cognite_toolkit._cdf_tk.utils import humanize_collection
+from cognite_toolkit._cdf_tk.utils.useful_types import T_ID
+
+Status: TypeAlias = Literal["success", "failed", "aborted", "pending"]
+
+AVAILABLE_STATUS = get_args(Status)
+
+
+class ProgressTracker(Generic[T_ID]):
+    """A thread-safe progress tracker for multiple items and steps.
+
+    For example, when migrating asset-centric Assets to CogniteAsset, we would like to track the progress of each Asset
+    through the steps of the migration. This is such that we can report back to the user which Assets failed, where
+    they failed, or if they succeeded.
+
+    Args:
+        steps (list[str]): A list of unique step names to track progress for.
+
+    Examples:
+        >>> tracker = ProgressTracker(steps=["downloadAsset", "convert", "uploadCogniteAssset"])
+        >>> tracker.set_progress(item_id=123, step="downloadAsset", status="success")
+        >>> tracker.set_progress(item_id=123, step="convert", status="failed")
+        >>> tracker.get_progress(item_id=123)
+        {'downloadAsset': 'success', 'convert': 'failed', 'uploadCogniteAssset': 'aborted'}
+
+    """
+
+    def __init__(self, steps: list[str]) -> None:
+        if len(steps) != len(set(steps)):
+            raise ValueError("Input `steps` must be unique.")
+        self._steps = list(steps)
+        self._step_to_idx: dict[str, int] = {step: i for i, step in enumerate(self._steps)}
+        self._progress: dict[T_ID, dict[str, Status]] = {}
+        self._lock = threading.Lock()
+
+    def _init_item(self, item_id: T_ID) -> None:
+        if item_id not in self._progress:
+            self._progress[item_id] = {step: "pending" for step in self._steps}
+
+    def set_progress(self, item_id: T_ID, step: str, status: Status) -> None:
+        """Set the progress of a specific step for a given item."""
+        if status not in AVAILABLE_STATUS:
+            raise ValueError(f"Status must be one of {humanize_collection(AVAILABLE_STATUS)}")
+        with self._lock:
+            self._init_item(item_id)
+            try:
+                idx = self._step_to_idx[step]
+            except KeyError as e:
+                raise ValueError(f"Step '{step}' not found in steps {humanize_collection(self._steps)}") from e
+            self._progress[item_id][step] = status
+            if status == "failed":
+                for s in self._steps[idx + 1 :]:
+                    self._progress[item_id][s] = "aborted"
+
+    @overload
+    def get_progress(self, item_id: T_ID) -> dict[str, Status]: ...
+
+    @overload
+    def get_progress(self, item_id: T_ID, step: str) -> Status: ...
+
+    def get_progress(self, item_id: T_ID, step: str | None = None) -> dict[str, Status] | Status:
+        """Get the progress of all steps or a specific step for a given item."""
+        with self._lock:
+            self._init_item(item_id)
+            if step is None:
+                return self._progress[item_id].copy()
+            if step not in self._step_to_idx:
+                raise ValueError(f"Step '{step}' not found in steps {humanize_collection(self._steps)}")
+            return self._progress[item_id][step]
+
+    def result(self) -> dict[T_ID, dict[str, Status]]:
+        """Get the progress of all items."""
+        with self._lock:
+            return {item_id: progress.copy() for item_id, progress in self._progress.items()}
+
+    def aggregate(self) -> dict[tuple[str, Status], int]:
+        """Aggregate the progress across all items and steps.
+
+        Returns:
+            dict[tuple[str, Status], int]: A dictionary with keys as (step, status) tuples and values as counts.
+        """
+        with self._lock:
+            return Counter((step, status) for progress in self._progress.values() for step, status in progress.items())

cognite_toolkit/_cdf_tk/utils/sentry_utils.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from cognite_toolkit._cdf_tk.exceptions import (
     ToolkitError,
@@ -9,9 +9,9 @@ if TYPE_CHECKING:
     from sentry_sdk.types import Hint as SentryHint
 
 
-def sentry_exception_filter(event: "SentryEvent", hint: "SentryHint") -> "Optional[SentryEvent]":
+def sentry_exception_filter(event: "SentryEvent", hint: "SentryHint") -> "SentryEvent | None":
     if "exc_info" in hint:
-        exc_type, exc_value, tb = hint["exc_info"]
+        _exc_type, exc_value, _tb = hint["exc_info"]
         # Returning None prevents the event from being sent to Sentry
         if isinstance(exc_value, ToolkitError):
             return None