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

cognite_toolkit/_cdf_tk/utils/cdf.py

@@ -1,6 +1,10 @@
 import sys
+import tempfile
+import time
 from collections.abc import Hashable, Iterator
 from dataclasses import dataclass
+from pathlib import Path
+from threading import RLock
 from typing import Any, Literal, overload
 from urllib.parse import urlparse
 
@@ -12,13 +16,19 @@ from cognite.client.data_classes import (
 from cognite.client.data_classes.data_modeling import Edge, Node, ViewId
 from cognite.client.data_classes.filters import SpaceFilter
 from cognite.client.exceptions import CogniteAPIError
+from cognite.client.utils.useful_types import SequenceNotStr
+from filelock import BaseFileLock, FileLock, Timeout
 from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient, ToolkitClientConfig
-from cognite_toolkit._cdf_tk.constants import ENV_VAR_PATTERN
+from cognite_toolkit._cdf_tk.client.data_classes.raw import RawTable
+from cognite_toolkit._cdf_tk.constants import ENV_VAR_PATTERN, MAX_ROW_ITERATION_RUN_QUERY, MAX_RUN_QUERY_FREQUENCY_MIN
 from cognite_toolkit._cdf_tk.exceptions import (
+    ToolkitError,
     ToolkitRequiredValueError,
+    ToolkitThrottledError,
     ToolkitTypeError,
+    ToolkitValueError,
 )
 from cognite_toolkit._cdf_tk.tk_warnings import (
     HighSeverityWarning,
@@ -26,6 +36,9 @@ from cognite_toolkit._cdf_tk.tk_warnings import (
 )
 from cognite_toolkit._cdf_tk.utils import humanize_collection
 
+from .file import sanitize_filename
+from .sql_parser import SQLParser
+
 if sys.version_info < (3, 11):
     from typing_extensions import Self
 else:
@@ -58,7 +71,7 @@ def try_find_error(credentials: OidcCredentials | ClientCredentials | None) -> s
 def iterate_instances(
     client: ToolkitClient,
     instance_type: Literal["node"] = "node",
-    space: str | None = None,
+    space: str | SequenceNotStr[str] | None = None,
     source: ViewId | None = None,
     console: Console | None = None,
 ) -> Iterator[Node]: ...
@@ -68,7 +81,7 @@ def iterate_instances(
 def iterate_instances(
     client: ToolkitClient,
     instance_type: Literal["edge"],
-    space: str | None = None,
+    space: str | SequenceNotStr[str] | None = None,
     source: ViewId | None = None,
     console: Console | None = None,
 ) -> Iterator[Edge]: ...
@@ -77,7 +90,7 @@ def iterate_instances(
 def iterate_instances(
     client: ToolkitClient,
     instance_type: Literal["node", "edge"] = "node",
-    space: str | None = None,
+    space: str | SequenceNotStr[str] | None = None,
     source: ViewId | None = None,
     console: Console | None = None,
 ) -> Iterator[Node] | Iterator[Edge]:
@@ -90,10 +103,14 @@ def iterate_instances(
     # By making the sort be on external_id, Postgres should pick the index that's on (project_id, space, external_id)
     # WHERE deleted_at IS NULL. In other words, avoiding soft deleted instances.
     body["sort"] = [
+        {
+            "property": [instance_type, "space"],
+            "direction": "ascending",
+        },
         {
             "property": [instance_type, "externalId"],
             "direction": "ascending",
-        }
+        },
     ]
     url = f"/api/{client._API_VERSION}/projects/{client.config.project}/models/instances/list"
     if space:
@@ -124,7 +141,7 @@ def iterate_instances(
 
 @overload
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -135,7 +152,7 @@ def read_auth(
 
 @overload
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -145,7 +162,7 @@ def read_auth(
 
 
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -172,6 +189,27 @@ def read_auth(
         return ClientCredentials(authentication["clientId"], authentication["clientSecret"])
 
 
+def get_transformation_sources(query: str) -> list[RawTable | str]:
+    """Search the SQL query for source tables."""
+    parser = SQLParser(query, operation="Lookup transformation source")
+    parser.parse()
+
+    tables: list[RawTable | str] = []
+    for table in parser.sources:
+        if table.schema == "_cdf":
+            tables.append(table.name)
+        else:
+            tables.append(RawTable(db_name=table.schema, table_name=table.name))
+    return tables
+
+
+def get_transformation_destination_columns(query: str) -> list[str]:
+    """Search the SQL query for destination columns."""
+    parser = SQLParser(query, operation="Lookup transformation destination columns")
+    parser.parse()
+    return parser.destination_columns
+
+
 def metadata_key_counts(
     client: ToolkitClient,
     resource: Literal["assets", "events", "files", "timeseries", "sequences"],
@@ -189,13 +227,11 @@ def metadata_key_counts(
     Returns:
         A dictionary with the metadata keys as keys and the counts as values.
     """
-    where_clause = ""
-    if data_sets is not None and hierarchies is not None:
-        where_clause = f"\n         WHERE dataSetId IN ({','.join(map(str, data_sets))}) AND rootId IN ({','.join(map(str, hierarchies))})"
-    elif data_sets is not None:
-        where_clause = f"\n         WHERE dataSetId IN ({','.join(map(str, data_sets))})"
-    elif hierarchies is not None:
-        where_clause = f"\n         WHERE rootId IN ({','.join(map(str, hierarchies))})"
+    if hierarchies and resource != "assets":
+        raise ToolkitValueError(
+            f"Hierarchies filtering for metadata keys are only supported for assets, but {resource} was provided."
+        )
+    where_clause = _create_where_clause(data_sets, hierarchies)
 
     query = f"""WITH meta AS (
          SELECT cast_to_strings(metadata) AS metadata_array
@@ -226,29 +262,56 @@ def metadata_key_counts(
     return [(item["key"], item["key_count"]) for item in results.results or []]
 
 
+def _create_where_clause(data_sets: list[int] | None, hierarchies: list[int] | None) -> str:
+    conditions = []
+    if data_sets:
+        if not all(isinstance(item, int) for item in data_sets):
+            raise ToolkitValueError("All items in data_sets must be integers for SQL filtering.")
+        conditions.append(f"dataSetId IN ({','.join(map(str, data_sets))})")
+    if hierarchies:
+        if not all(isinstance(item, int) for item in hierarchies):
+            raise ToolkitValueError("All items in hierarchies must be integers for SQL filtering.")
+        conditions.append(f"rootId IN ({','.join(map(str, hierarchies))})")
+
+    if not conditions:
+        return ""
+    return f"\n         WHERE {' AND '.join(conditions)}"
+
+
 def label_count(
-    client: ToolkitClient, resource: Literal["assets", "events", "files", "timeseries", "sequences"]
-) -> list[dict[str, int | str]]:
+    client: ToolkitClient,
+    resource: Literal["assets", "events", "files", "timeseries", "sequences"],
+    data_sets: list[int] | None = None,
+    hierarchies: list[int] | None = None,
+) -> list[tuple[str, int]]:
     """Get the label counts for a given resource.
 
     Args:
         client: ToolkitClient instance
         resource: The resource to get the label counts for. Can be one of "assets", "events", "files", "timeseries", or "sequences".
+        data_sets: A list of data set IDs to filter by. If None, no filtering is applied.
+        hierarchies: A list of hierarchy IDs to filter by. If None, no filtering is applied.
 
     Returns:
-        A dictionary with the labels as keys and the counts as values.
+        A list of tuples with the label and its count.
     """
+    if hierarchies and resource != "assets":
+        raise ToolkitValueError(
+            f"Hierarchies filtering for labels are only supported for assets, but {resource} was provided."
+        )
+    where_clause = _create_where_clause(data_sets, hierarchies)
+
     query = f"""WITH labels as (SELECT explode(labels) AS label
-	FROM _cdf.{resource}
+	FROM _cdf.{resource}{where_clause}
   )
 SELECT label, COUNT(label) as label_count
 FROM labels
 GROUP BY label
 ORDER BY label_count DESC;
 """
-    results = client.transformations.preview(query, convert_to_string=False, limit=1000)
+    results = client.transformations.preview(query, convert_to_string=False, limit=None, source_limit=None)
     # We know from the SQL that the result is a list of dictionaries with string keys and int values.
-    return results.results or []
+    return [(item["label"], item["label_count"]) for item in results.results or []]
 
 
 @dataclass
@@ -318,3 +381,110 @@ FROM
     if results.results:
         return int(results.results[0]["labelCount"])
     return 0
+
+
+@dataclass
+class ThrottlerState:
+    _FILENAME = "tk-last-raw_count-{project}.txt"
+
+    lock: BaseFileLock
+    project: str
+    _allowed_this_run: bool | None = None
+
+    @classmethod
+    def get(cls, project: str) -> "ThrottlerState":
+        if project in _STATE_BY_PROJECT:
+            return _STATE_BY_PROJECT[project]
+
+        with _STATE_LOCK:
+            # Re-check after acquiring the lock to handle the race condition.
+            if project in _STATE_BY_PROJECT:
+                return _STATE_BY_PROJECT[project]
+            filepath = cls._filepath(project)
+            # filelock needs a string path
+            lock = FileLock(str(filepath.with_suffix(".lock")), timeout=10.0)
+            state = cls(lock=lock, project=project)
+            _STATE_BY_PROJECT[project] = state
+            return state
+
+    @classmethod
+    def _filepath(cls, project: str) -> Path:
+        # Enure the project name is a valid string for hashing.
+        project = sanitize_filename(project)
+        filename = cls._FILENAME.format(project=project)
+        return Path(tempfile.gettempdir()) / filename
+
+    def throttle(self) -> None:
+        """Checks if the function can be run, and if so, updates the last call timestamp.
+
+        This is an atomic operation protected by a file lock.
+
+        Raises:
+            ToolkitThrottledError: If the function has been called too recently.
+        """
+        try:
+            with self.lock:
+                now = time.time()
+                filepath = self._filepath(self.project)
+                if self._allowed_this_run is None:
+                    last_call_epoch = 0.0
+                    if filepath.exists():
+                        try:
+                            last_call_epoch = float(filepath.read_text(encoding="utf-8"))
+                        except (ValueError, FileNotFoundError):
+                            # File is corrupt or was deleted after check. Allow to run and overwrite.
+                            pass
+
+                    to_wait = (MAX_RUN_QUERY_FREQUENCY_MIN * 60) - (now - last_call_epoch)
+                    if to_wait > 0:
+                        raise ToolkitThrottledError(
+                            f"Row count is limited to once every {MAX_RUN_QUERY_FREQUENCY_MIN} minutes. "
+                            f"Please wait {to_wait:.2f} seconds before calling again.",
+                            to_wait,
+                        )
+                    self._allowed_this_run = True
+
+                # We are allowed to run, so we update the timestamp.
+                filepath.write_text(str(now), encoding="utf-8")
+        except Timeout:
+            raise ToolkitError(
+                "Could not acquire lock for throttling raw row count. "
+                "Another toolkit process might be running. Please try again."
+            )
+
+
+_STATE_LOCK = RLock()
+_STATE_BY_PROJECT: dict[str, ThrottlerState] = {}
+
+
+def raw_row_count(client: ToolkitClient, raw_table_id: RawTable, max_count: int = MAX_ROW_ITERATION_RUN_QUERY) -> int:
+    """Get the number of rows in a raw table.
+
+    Args:
+        client: ToolkitClient instance
+        raw_table_id: The ID of the raw table to count rows in.
+        max_count: The maximum number of rows to count. If the table has more rows than this, the count will be capped.
+
+    Returns:
+        The number of rows in the raw table.
+
+    Raises:
+        ToolkitThrottledError: If the row count is limited to once every MAX_RUN_QUERY_FREQUENCY_MIN minutes.
+        ValueError: If max_count is not between 0 and MAX_ROW_ITERATION_RUN_QUERY (inclusive).
+    """
+    ThrottlerState.get(client.config.project).throttle()
+
+    if not 0 <= max_count <= MAX_ROW_ITERATION_RUN_QUERY:
+        raise ValueError(f"max_count must be between 0 and {MAX_ROW_ITERATION_RUN_QUERY} (inclusive).")
+
+    query = f"""SELECT COUNT(key) AS row_count
+FROM (
+    SELECT key
+    FROM `{raw_table_id.db_name}`.`{raw_table_id.table_name}`
+    LIMIT {max_count}
+) AS limited_keys"""
+
+    results = client.transformations.preview(query, convert_to_string=False, limit=None, source_limit=None)
+    if results.results:
+        return int(results.results[0]["row_count"])
+    return 0

cognite_toolkit/_cdf_tk/utils/cli_args.py

@@ -0,0 +1,34 @@
+from cognite.client.data_classes.data_modeling import ViewId
+
+from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
+
+
+def parse_view_str(view_str: str) -> ViewId:
+    """Parse a view string into a ViewId.
+
+    Args:
+        view_str (str): The view string to parse.
+
+    Returns:
+        ViewId: The parsed ViewId.
+
+    Raises:
+        ToolkitValueError: If the view string is not in a valid format.
+
+    >>> parse_view_str("my_space:my_view/v1")
+    ViewId(space='my_space', external_id='my_view', version='v1')
+    """
+    space_and_rest = view_str.split(":", 1)
+    if len(space_and_rest) != 2 or not space_and_rest[0] or not space_and_rest[1]:
+        raise ToolkitValueError(
+            f"Invalid view string format: '{view_str}'. Expected format 'space:externalId/version'."
+        )
+    space = space_and_rest[0]
+    external_id_and_version = space_and_rest[1].rsplit("/", 1)
+    if len(external_id_and_version) != 2 or not external_id_and_version[0] or not external_id_and_version[1]:
+        raise ToolkitValueError(
+            f"Invalid view string format: '{view_str}'. Expected format 'space:externalId/version'."
+        )
+    external_id = external_id_and_version[0]
+    version = external_id_and_version[1]
+    return ViewId(space=space, external_id=external_id, version=version)

cognite_toolkit/_cdf_tk/utils/collection.py

@@ -1,7 +1,7 @@
 import difflib
-from collections.abc import Collection, Iterable, Iterator
+from collections.abc import Collection, Iterable, Iterator, Sequence, Set
 from itertools import islice
-from typing import Any
+from typing import Any, TypeVar
 
 from cognite_toolkit._cdf_tk.utils.file import yaml_safe_dump
 
@@ -18,6 +18,47 @@ def flatten_dict(dct: dict[str, Any]) -> dict[tuple[str, ...], Any]:
     return items
 
 
+def flatten_dict_json_path(dct: dict[str, Any], keep_structured: Set[str] | None = None) -> dict[str, Any]:
+    """Flatten a dictionary to a dictionary with JSON path keys.
+
+    Empty keys are ignored.
+
+    Args:
+        dct: The dictionary to flatten.
+        keep_structured: A set of keys to keep structured (not flatten). If a key is in this set,
+            it will not be flattened, and the value will be kept as is. This only applies to top-level keys.
+
+    Returns:
+        A dictionary with JSON path keys.
+    """
+
+    return _flatten(dct, keep_structured=keep_structured or set())
+
+
+def _flatten(obj: Any, keep_structured: Set[str], path: str = "") -> dict[str, Any]:
+    items: dict[str, Any] = {}
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if not key:
+                continue
+            current_path = f"{path}.{key}" if path else key
+            if key in keep_structured:
+                items[current_path] = value
+            else:
+                items.update(_flatten(value, set(), current_path))
+    elif isinstance(obj, list):
+        for i, value in enumerate(obj):
+            current_path = f"{path}[{i}]"
+            if isinstance(value, (dict, list)):
+                items.update(_flatten(value, set(), current_path))
+            else:
+                items[current_path] = value
+    else:
+        items[path] = obj
+
+    return items
+
+
 def to_diff(a: dict[str, Any], b: dict[str, Any]) -> Iterator[str]:
     a_str = yaml_safe_dump(a, sort_keys=True)
     b_str = yaml_safe_dump(b, sort_keys=True)
@@ -48,3 +89,13 @@ def chunker(iterable: Iterable[Any], size: int) -> Iterator[list[Any]]:
     iterator = iter(iterable)
     while chunk := list(islice(iterator, size)):
         yield chunk
+
+
+T_Sequence = TypeVar("T_Sequence", bound=Sequence)
+
+
+def chunker_sequence(sequence: T_Sequence, size: int) -> Iterator[T_Sequence]:
+    """Yield successive n-sized chunks from sequence."""
+    for i in range(0, len(sequence), size):
+        # MyPy does not expect sequence[i : i + size] to be of type T_Sequence
+        yield sequence[i : i + size]  # type: ignore[misc]