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

cognite_toolkit/_cdf_tk/commands/modules.py

@@ -1,12 +1,15 @@
-from __future__ import annotations
-
 import shutil
+import sys
+import tempfile
+import zipfile
 from collections import Counter
-from importlib import resources
+from hashlib import sha256
 from pathlib import Path
-from typing import Any, Literal, Optional
+from types import TracebackType
+from typing import Any, Literal
 
 import questionary
+import requests
 import typer
 from packaging.version import Version
 from packaging.version import parse as parse_version
@@ -14,13 +17,12 @@ from rich import print
 from rich.markdown import Markdown
 from rich.padding import Padding
 from rich.panel import Panel
-from rich.progress import track
+from rich.progress import Progress, track
 from rich.rule import Rule
 from rich.table import Table
 from rich.tree import Tree
 
-import cognite_toolkit
-from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
+from cognite_toolkit._cdf_tk.cdf_toml import CDFToml, Library
 from cognite_toolkit._cdf_tk.commands import _cli_commands as CLICommands
 from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
 from cognite_toolkit._cdf_tk.commands._changes import (
@@ -32,9 +34,10 @@ from cognite_toolkit._cdf_tk.commands._changes import (
     UpdateDockerImageVersion,
     UpdateModuleVersion,
 )
+from cognite_toolkit._cdf_tk.commands._questionary_style import custom_style_fancy
 from cognite_toolkit._cdf_tk.constants import (
-    BUILTIN_MODULES,
     MODULES,
+    RESOURCES_PATH,
     SUPPORT_MODULE_UPGRADE_FROM_VERSION,
     EnvType,
 )
@@ -47,29 +50,20 @@ from cognite_toolkit._cdf_tk.data_classes import (
     Package,
     Packages,
 )
-from cognite_toolkit._cdf_tk.exceptions import ToolkitRequiredValueError, ToolkitValueError
+from cognite_toolkit._cdf_tk.exceptions import ToolkitError, ToolkitRequiredValueError, ToolkitValueError
 from cognite_toolkit._cdf_tk.hints import verify_module_directory
 from cognite_toolkit._cdf_tk.tk_warnings import MediumSeverityWarning
+from cognite_toolkit._cdf_tk.tk_warnings.other import HighSeverityWarning
 from cognite_toolkit._cdf_tk.utils import humanize_collection, read_yaml_file
 from cognite_toolkit._cdf_tk.utils.file import safe_read, safe_rmtree, safe_write, yaml_safe_dump
 from cognite_toolkit._cdf_tk.utils.modules import module_directory_from_path
 from cognite_toolkit._cdf_tk.utils.repository import FileDownloader
 from cognite_toolkit._version import __version__
 
-custom_style_fancy = questionary.Style(
-    [
-        ("qmark", "fg:#673ab7"),  # token in front of the question
-        ("question", "bold"),  # question text
-        ("answer", "fg:#f44336 bold"),  # submitted answer text behind the question
-        ("pointer", "fg:#673ab7 bold"),  # pointer used in select and checkbox prompts
-        ("highlighted", "fg:#673ab7 bold"),  # pointed-at choice in select and checkbox prompts
-        ("selected", "fg:#673ab7"),  # style for a selected item of a checkbox
-        ("separator", "fg:#cc5454"),  # separator in lists
-        ("instruction", ""),  # user instructions for select, rawselect, checkbox
-        ("text", ""),  # plain text
-        ("disabled", "fg:#858585 italic"),  # disabled choices for select and checkbox prompts
-    ]
-)
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
 
 INDENT = "  "
 POINTER = INDENT + "▶"
@@ -82,9 +76,40 @@ _FILE_DOWNLOADERS_BY_TYPE: dict[str, type[FileDownloader]] = {
 
 
 class ModulesCommand(ToolkitCommand):
-    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False):
+    def __init__(
+        self,
+        print_warning: bool = True,
+        skip_tracking: bool = False,
+        silent: bool = False,
+        temp_dir_suffix: str | None = None,
+        module_source_dir: Path | None = None,
+    ):
         super().__init__(print_warning, skip_tracking, silent)
-        self._builtin_modules_path = Path(resources.files(cognite_toolkit.__name__)) / BUILTIN_MODULES  # type: ignore [arg-type]
+        self._module_source_dir: Path | None = module_source_dir
+        # Use suffix to make temp directory unique (useful for parallel test execution)
+        modules_dir_name = f"{MODULES}.{temp_dir_suffix}" if temp_dir_suffix else MODULES
+        self._temp_download_dir = Path(tempfile.gettempdir()) / modules_dir_name
+        if not self._temp_download_dir.exists():
+            self._temp_download_dir.mkdir(parents=True, exist_ok=True)
+
+    def __enter__(self) -> Self:
+        """
+        Context manager to ensure the temporary download directory is cleaned up after use. It requires the command to be used in a `with` block.
+        """
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,  # Type of the exception
+        exc_value: BaseException | None,  # Exception instance
+        traceback: TracebackType | None,  # Traceback object
+    ) -> None:
+        """
+        Clean up the temporary download directory.
+        """
+
+        if self._temp_download_dir.exists():
+            safe_rmtree(self._temp_download_dir)
 
     @classmethod
     def _create_tree(cls, item: Packages) -> Tree:
@@ -113,25 +138,32 @@ class ModulesCommand(ToolkitCommand):
         organization_dir: Path,
         selected_packages: Packages,
         environments: list[EnvType],
+        modules_source_path: Path,
         mode: Literal["new", "clean", "update"] | None,
         download_data: bool = False,
     ) -> None:
-        modules_root_dir = organization_dir / MODULES
-        if mode == "clean" and modules_root_dir.is_dir():
+        modules_target_root_dir = organization_dir / MODULES
+        if mode == "clean" and modules_target_root_dir.is_dir():
             print(f"{INDENT}[yellow]Clearing directory[/]")
-            safe_rmtree(modules_root_dir)
+            safe_rmtree(modules_target_root_dir)
 
-        modules_root_dir.mkdir(parents=True, exist_ok=True)
+        modules_target_root_dir.mkdir(parents=True, exist_ok=True)
 
         seen_modules: set[Path] = set()
         selected_paths: set[Path] = set()
         downloader_by_repo: dict[str, FileDownloader] = {}
 
         extra_resources: set[Path] = set()
+
         for package_name, package in selected_packages.items():
             print(f"{INDENT}[{'yellow' if mode == 'clean' else 'green'}]Creating {package_name}[/]")
 
             for module in package.modules:
+                if module.module_id:
+                    self._additional_tracking_info.installed_module_ids.add(module.module_id)
+                if module.package_id:
+                    self._additional_tracking_info.installed_package_ids.add(module.package_id)
+
                 if module.dir in seen_modules:
                     # A module can be part of multiple packages
                     continue
@@ -144,7 +176,7 @@ class ModulesCommand(ToolkitCommand):
                     extra_resources.update(module.definition.extra_resources)
 
                 print(f"{INDENT * 2}[{'yellow' if mode == 'clean' else 'green'}]Creating module {module.name}[/]")
-                target_dir = modules_root_dir / module.relative_path
+                target_dir = modules_target_root_dir / module.relative_path
                 if Path(target_dir).exists() and mode == "update":
                     if questionary.confirm(
                         f"{INDENT}Module {module.name} already exists in folder {target_dir}. Would you like to overwrite?",
@@ -180,8 +212,8 @@ class ModulesCommand(ToolkitCommand):
             created_by_module: dict[Path, int] = Counter()
             for extra in extra_resources:
                 module_dir = module_directory_from_path(extra)
-                extra_full_path = self._builtin_modules_path / extra
-                target_path = modules_root_dir / extra
+                extra_full_path = modules_source_path / extra
+                target_path = modules_target_root_dir / extra
                 if target_path.exists():
                     # Assume that the user has already created this shared resource
                     continue
@@ -206,7 +238,7 @@ class ModulesCommand(ToolkitCommand):
             if mode == "update":
                 config_init = InitConfigYAML.load_existing(
                     safe_read(Path(organization_dir) / f"config.{environment}.yaml"), environment
-                ).load_defaults(self._builtin_modules_path, selected_paths)
+                ).load_defaults(modules_source_path, selected_paths)
             else:
                 ignore_variable_patterns: list[tuple[str, ...]] | None = None
                 if len(selected_packages) == 1 and "quickstart" in selected_packages:
@@ -229,7 +261,7 @@ class ModulesCommand(ToolkitCommand):
                         validation_type=environment,
                         selected=[f"{MODULES}/"],
                     )
-                ).load_defaults(self._builtin_modules_path, selected_paths, ignore_variable_patterns)
+                ).load_defaults(modules_source_path, selected_paths, ignore_variable_patterns)
 
                 if len(selected_packages) == 1 and "quickstart" in selected_packages:
                     config_init.lift()
@@ -239,54 +271,50 @@ class ModulesCommand(ToolkitCommand):
             )
             safe_write(Path(organization_dir) / f"config.{environment}.yaml", config_init.dump_yaml_with_comments())
 
-        cdf_toml_content = self.create_cdf_toml(organization_dir, environments[0] if environments else "dev")
-
-        destination = Path.cwd() / CDFToml.file_name
-        if destination.exists():
-            print(f"{INDENT}[yellow]cdf.toml file already exists. Skipping creation.")
-        else:
-            destination.write_text(cdf_toml_content, encoding="utf-8")
-
-    def create_cdf_toml(self, organization_dir: Path, env: EnvType = "dev") -> str:
-        cdf_toml_content = safe_read(self._builtin_modules_path / CDFToml.file_name)
-        if organization_dir != Path.cwd():
-            cdf_toml_content = cdf_toml_content.replace(
-                "#<PLACEHOLDER>",
-                f'''
-default_organization_dir = "{organization_dir.name}"''',
-            )
-        else:
-            cdf_toml_content = cdf_toml_content.replace("#<PLACEHOLDER>", "")
-        cdf_toml_content = cdf_toml_content.replace("<DEFAULT_ENV_PLACEHOLDER>", env)
-        return cdf_toml_content
+        CDFToml.write(organization_dir, environments[0] if environments else "dev")
 
     def init(
         self,
-        organization_dir: Optional[Path] = None,
+        organization_dir: Path | None = None,
         select_all: bool = False,
         clean: bool = False,
         user_select: str | None = None,
         user_environments: list[str] | None = None,
         user_download_data: bool | None = None,
+        library_url: str | None = None,
+        library_checksum: str | None = None,
     ) -> None:
         if not organization_dir:
-            new_line = "\n    "
-            message = (
-                f"Which directory would you like to create templates in? (default: current directory){new_line}"
-                f"HINT Use an organization directory if you use the repository for more than Toolkit. "
-                f"If not, use the current (repository root) directory '.':"
-            )
-            organization_dir_raw = questionary.text(message=message, default="").ask()
-            organization_dir = Path(organization_dir_raw.strip())
+            organization_dir = ModulesCommand._prompt_organization_dir()
 
         modules_root_dir = organization_dir / MODULES
-        packages = Packages().load(self._builtin_modules_path)
+
+        # Determine which library to use (if any)
+        library: Library | None = None
+        if library_url:
+            if not library_checksum:
+                raise ToolkitRequiredValueError(
+                    "The '--library-checksum' is required when '--library-url' is provided."
+                )
+            library = Library(url=library_url, checksum=library_checksum)
+        elif not (organization_dir / CDFToml.file_name).exists():
+            # Load default library from resources when cdf.toml doesn't exist
+            default_cdf_toml = CDFToml.load(cwd=RESOURCES_PATH, use_singleton=False)
+            library = default_cdf_toml.libraries.get("toolkit-data")
+            if library is None:
+                raise ToolkitError("Default cdf.toml in resources is missing 'toolkit-data' library configuration.")
+
+        packages, modules_source_path = self._get_available_packages(library)
 
         if select_all:
             print(Panel("Instantiating all available modules"))
             mode = self._verify_clean(modules_root_dir, clean)
             self._create(
-                organization_dir=organization_dir, selected_packages=packages, environments=["dev", "prod"], mode=mode
+                organization_dir=organization_dir,
+                selected_packages=packages,
+                environments=["dev", "prod"],
+                mode=mode,
+                modules_source_path=modules_source_path,
             )
             return
         is_interactive = user_select is not None
@@ -324,7 +352,13 @@ default_organization_dir = "{organization_dir.name}"''',
                 mode = self._verify_clean(bootcamp_org / MODULES, clean)
             # We only ask to verify if the user has not already selected overwrite in the _verify_clean method
             if mode == "clean" or questionary.confirm("Would you like to continue with creation?", default=True).ask():
-                self._create(bootcamp_org, selected, ["test"], mode)
+                self._create(
+                    organization_dir=bootcamp_org,
+                    selected_packages=selected,
+                    environments=["test"],
+                    mode=mode,
+                    modules_source_path=modules_source_path,
+                )
             raise typer.Exit()
 
         if (
@@ -354,7 +388,14 @@ default_organization_dir = "{organization_dir.name}"''',
             download_data = self._get_download_data(selected)
         else:
             download_data = user_download_data
-        self._create(organization_dir, selected, environments, mode, download_data)
+        self._create(
+            organization_dir=organization_dir,
+            selected_packages=selected,
+            environments=environments,
+            mode=mode,
+            download_data=download_data,
+            modules_source_path=modules_source_path,
+        )
 
         print(
             Panel(
@@ -372,6 +413,21 @@ default_organization_dir = "{organization_dir.name}"''',
         if not is_interactive:
             raise typer.Exit()
 
+    @staticmethod
+    def _prompt_organization_dir() -> Path:
+        print(
+            "\n".join(
+                [
+                    "Enter an optional subdirectory for Toolkit content.",
+                    "    [bold yellow]HINT[/bold yellow] Use an organization directory if you use the repository for more than Toolkit.",
+                    "    [dim]The current directory is the default.[/dim]",
+                ]
+            )
+        )
+
+        organization_dir_raw = questionary.text(message="", default="").ask()
+        return Path(organization_dir_raw.strip())
+
     @staticmethod
     def _get_download_data(selected: Packages) -> bool:
         example_data = {
@@ -394,7 +450,7 @@ default_organization_dir = "{organization_dir.name}"''',
         choices = sorted(
             [
                 questionary.Choice(
-                    title=module.title,
+                    title=module.title or module.name,
                     value=module,
                     checked=module.name in dependencies or module.is_selected_by_default,
                     disabled="required" if module.name in dependencies else None,
@@ -680,9 +736,201 @@ default_organization_dir = "{organization_dir.name}"''',
             build_env = default.environment.validation_type
 
         existing_module_names = [module.name for module in ModuleResources(organization_dir, build_env).list()]
-        available_packages = Packages().load(self._builtin_modules_path)
-
+        available_packages, modules_source_path = self._get_available_packages()
         added_packages = self._select_packages(available_packages, existing_module_names)
 
         download_data = self._get_download_data(added_packages)
-        self._create(organization_dir, added_packages, environments, "update", download_data)
+        self._create(
+            organization_dir=organization_dir,
+            selected_packages=added_packages,
+            environments=environments,
+            mode="update",
+            download_data=download_data,
+            modules_source_path=modules_source_path,
+        )
+
+    def _get_available_packages(self, user_library: Library | None = None) -> tuple[Packages, Path]:
+        """
+        Returns a list of available packages, either from the CDF TOML file or from external libraries if the feature flag is enabled.
+        """
+
+        cdf_toml = CDFToml.load()
+
+        if self._module_source_dir is None:
+            libraries = {"userdefined": user_library} if user_library else cdf_toml.libraries
+
+            for library_name, library in libraries.items():
+                try:
+                    if library_name:
+                        self._additional_tracking_info.downloaded_library_ids.add(library_name)
+
+                    print(f"[green]Adding library {library_name} from {library.url}[/]")
+                    # Extract filename from URL, fallback to library_name.zip if no filename found
+                    from urllib.parse import urlparse
+
+                    url_path = urlparse(library.url).path
+                    filename = (
+                        url_path.split("/")[-1] if url_path and url_path.split("/")[-1] else f"{library_name}.zip"
+                    )
+                    file_path = self._temp_download_dir / filename
+                    self._download(library.url, file_path)
+                    self._validate_checksum(library.checksum, file_path)
+                    self._unpack(file_path)
+                    packages = Packages().load(file_path.parent)
+                    if packages.warnings:
+                        for warning in packages.warnings:
+                            self.warn(warning)
+                    self._validate_packages(packages, f"library {library_name}")
+
+                    # Track deployment pack download for each package and module
+                    for package in packages.values():
+                        if package.id:
+                            self._additional_tracking_info.downloaded_package_ids.add(package.id)
+
+                        for module in package.modules:
+                            if module.module_id:
+                                self._additional_tracking_info.downloaded_module_ids.add(module.module_id)
+
+                    return packages, file_path.parent
+                except Exception as e:
+                    if isinstance(e, ToolkitError):
+                        raise e
+                    else:
+                        raise ToolkitError(
+                            f"An unexpected error occurred during downloading {library.url} to {file_path}: {e}"
+                        ) from e
+
+                    raise ToolkitError(f"Failed to add library {library_name}, {e}")
+            # If no libraries are specified or the flag is not enabled, raise an error
+            raise ValueError("No valid libraries found.")
+        else:
+            if user_library:
+                self.warn(
+                    MediumSeverityWarning(
+                        "External library provided but not enabled in cdf.toml. Please enable the feature flag and try again."
+                    )
+                )
+
+            if self._module_source_dir is None:
+                self.warn(
+                    MediumSeverityWarning("No external libraries and no local module source directory specified.")
+                )
+                return Packages(), Path(".")
+            packages = Packages.load(self._module_source_dir)
+            self._validate_packages(packages, self._module_source_dir.name)
+            return packages, self._module_source_dir
+
+    def _validate_packages(self, packages: Packages, source_name: str) -> None:
+        """
+        Validates that packages are properly configured and have accessible modules.
+        Raises descriptive errors for common misconfigurations.
+        """
+        if not packages:
+            raise ToolkitError(
+                f"No packages found in {source_name}. "
+                "This may be due to a misconfigured packages.toml file or missing module directories."
+            )
+
+        # Check for packages with no modules
+        empty_packages = [pkg.name for pkg in packages.values() if not pkg.modules]
+        if empty_packages:
+            self.warn(
+                MediumSeverityWarning(
+                    f"The following packages in {source_name} have no modules: {', '.join(empty_packages)}. "
+                    "This may indicate a misconfiguration in packages.toml."
+                )
+            )
+
+        # Check for missing module directories
+        missing_paths = []
+        for package in packages.values():
+            for module in package.modules:
+                if not module.dir.exists():
+                    missing_paths.append(f"{module.name} ({module.dir})")
+
+        if missing_paths:
+            raise ToolkitError(
+                f"The following module paths in {source_name} do not exist:\n"
+                f"{chr(10).join(f'  - {path}' for path in missing_paths)}\n"
+                "This may be due to:\n"
+                "- Incorrect paths in packages.toml\n"
+                "- Missing module directories\n"
+                "- Corrupted or incomplete module installation"
+            )
+
+    def _download(self, url: str, file_path: Path) -> None:
+        """
+        Downloads a file from a URL to the specified output path.
+        If the file already exists, it deletes it before downloading.
+        """
+        try:
+            # normally the context manager exit should handle this, but if the process is killed the file will not be deleted
+            if file_path.exists():
+                file_path.unlink()
+
+            response = requests.get(url, stream=True)
+            response.raise_for_status()  # Raise an exception for HTTP errors
+
+            total_size = int(response.headers.get("content-length", 0))
+
+            with Progress() as progress:
+                task = progress.add_task("Download", total=total_size)
+                with open(file_path, "wb") as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        progress.update(task, advance=len(chunk))
+
+        except requests.exceptions.RequestException as e:
+            raise ToolkitError(f"Error downloading file from {url}: {e}") from e
+
+    def _validate_checksum(self, checksum: str, file_path: Path) -> None:
+        """
+        Compares the checksum of the downloaded file with the expected checksum.
+        """
+
+        if checksum.lower().startswith("sha256:"):
+            checksum = checksum[7:]
+        else:
+            raise ToolkitValueError(f"Unsupported checksum format: {checksum}. Expected 'sha256:' prefix")
+
+        chunk_size: int = 8192
+        sha256_hash = sha256()
+        try:
+            with open(file_path, "rb") as f:
+                # Read the file in chunks to handle large files efficiently
+                for chunk in iter(lambda: f.read(chunk_size), b""):
+                    sha256_hash.update(chunk)
+            calculated = sha256_hash.hexdigest()
+        except OSError as e:
+            raise ToolkitError(f"Failed to calculate checksum for {file_path}: {e}") from e
+        except Exception as e:
+            raise ToolkitError(f"Unexpected error during checksum calculation for {file_path}: {e}") from e
+
+        if calculated != checksum:
+            self.warn(
+                HighSeverityWarning(
+                    f"The provided checksum sha256:{checksum} does not match downloaded file hash sha256:{calculated}.\n"
+                    "Please verify the checksum with the source and update cdf.toml if needed.\n"
+                    "This may indicate that the package content has changed."
+                )
+            )
+        else:
+            print("[green]✓ Checksum verified[/green]")
+
+    def _unpack(self, file_path: Path) -> None:
+        """
+        Unzips the downloaded file to the specified output path.
+        If the file is not a zip file, it raises an error.
+        """
+        total_size = file_path.stat().st_size if file_path.exists() else 0
+
+        try:
+            with Progress() as progress:
+                unzip = progress.add_task("Unzipping", total=total_size)
+                with zipfile.ZipFile(file_path, "r") as zip_ref:
+                    zip_ref.extractall(file_path.parent)
+                    progress.update(unzip, advance=total_size)
+        except zipfile.BadZipFile as e:
+            raise ToolkitError(f"Error unpacking zip file {file_path}: {e}") from e
+        except Exception as e:
+            raise ToolkitError(f"An unexpected error occurred while unpacking {file_path}: {e}") from e