industrial-model 0.1.0__py3-none-any.whl

industrial_model/cognite_adapters/schemas.py

@@ -0,0 +1,121 @@
+import inspect
+from collections import defaultdict
+from types import UnionType
+from typing import (
+    Any,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from pydantic import BaseModel
+
+from industrial_model.constants import NESTED_SEP
+
+TBaseModel = TypeVar("TBaseModel", bound=BaseModel)
+
+
+def get_schema_properties(
+    cls: type[TBaseModel],
+    nested_separator: str = NESTED_SEP,
+    prefix: str | None = None,
+) -> list[str]:
+    data = _get_type_properties(cls, defaultdict(lambda: 0)) or {}
+    keys = _flatten_dict_keys(data, None, nested_separator)
+    if not prefix:
+        return keys
+
+    return [f"{prefix}{nested_separator}{key}" for key in keys]
+
+
+def _get_type_properties(
+    cls: type[BaseModel], visited_count: defaultdict[type, int]
+) -> dict[str, Any] | None:
+    if visited_count[cls] > 2:
+        return None
+    visited_count[cls] += visited_count[cls] + 1
+
+    hints = get_type_hints(cls)
+    origins = {
+        key: _get_field_type(type_hint, visited_count)
+        for key, type_hint in hints.items()
+    }
+
+    return {
+        field_info.alias or key: origins[key][1]
+        for key, field_info in cls.model_fields.items()
+        if key in origins
+    }
+
+
+def _get_field_type(
+    type_hint: type, visited_count: defaultdict[type, int]
+) -> tuple[bool, dict[str, Any] | None]:
+    should_iter = _type_is_list_or_union(type_hint)
+
+    if not should_iter:
+        return _get_field_relations(
+            [_cast_base_model(type_hint)], visited_count
+        )
+
+    entries: list[type[BaseModel] | None] = []
+    for arg in get_args(type_hint):
+        if _type_is_list_or_union(arg):
+            return _get_field_type(arg, visited_count)
+        entries.append(_cast_base_model(arg))
+
+    return _get_field_relations(entries, visited_count)
+
+
+def _get_field_relations(
+    entries: list[type[TBaseModel] | None],
+    visited_count: defaultdict[type, int],
+) -> tuple[bool, dict[str, Any] | None]:
+    entry_type = next((type_ for type_ in entries if type_ is not None), None)
+
+    if not entry_type:
+        return False, None
+
+    properties = _get_type_properties(entry_type, visited_count)
+
+    return True, properties
+
+
+def _type_is_list_or_union(entry: type) -> bool:
+    origin = get_origin(entry)
+    is_union = origin in (Union, UnionType)
+    is_list = origin in (list, list)
+
+    return is_union or is_list
+
+
+def _cast_base_model(entry: type) -> type[TBaseModel] | None:
+    is_base_model = (
+        entry is not type(None)
+        and inspect.isclass(entry) is True
+        and issubclass(entry, BaseModel)
+    )
+    return entry if is_base_model else None
+
+
+def _flatten_dict_keys(
+    data: dict[str, Any], parent_key: str | None, nested_separator: str
+) -> list[str]:
+    paths: set[str] = set()
+    for key, value in data.items():
+        full_key = (
+            f"{parent_key}{nested_separator}{key}" if parent_key else key
+        )
+        paths.add(full_key)
+        if isinstance(value, dict) and value:
+            paths.update(_flatten_dict_keys(value, full_key, nested_separator))
+        elif isinstance(value, str):
+            paths.add(f"{full_key}{nested_separator}{value}")
+        elif isinstance(value, list | set):
+            paths.update(
+                [f"{full_key}{nested_separator}{item}" for item in value]
+            )
+
+    return list(paths)

industrial_model/cognite_adapters/sort_mapper.py

@@ -0,0 +1,21 @@
+from cognite.client.data_classes.data_modeling import InstanceSort, View
+
+from industrial_model.cognite_adapters.utils import get_property_ref
+from industrial_model.constants import SORT_DIRECTION
+from industrial_model.statements.expressions import Column
+
+
+class SortMapper:
+    def map(
+        self,
+        sort_clauses: list[tuple[Column, SORT_DIRECTION]],
+        root_view: View,
+    ) -> list[InstanceSort]:
+        return [
+            InstanceSort(
+                property=get_property_ref(column.property, root_view),
+                direction=direction,
+                nulls_first=direction == "descending",
+            )
+            for column, direction in sort_clauses
+        ]

industrial_model/cognite_adapters/utils.py

@@ -0,0 +1,33 @@
+from typing import Literal
+
+from cognite.client.data_classes.data_modeling import (
+    View,
+    ViewId,
+)
+
+from industrial_model.models import TViewInstance
+
+NODE_PROPERTIES = {"externalId", "space", "createdTime", "deletedTime"}
+INTANCE_TYPE = Literal["node", "edge"]
+
+
+def get_property_ref(
+    property: str, view: View | ViewId, instance_type: INTANCE_TYPE = "node"
+) -> tuple[str, str, str] | tuple[str, str]:
+    return (
+        view.as_property_ref(property)
+        if property not in NODE_PROPERTIES
+        else (instance_type, property)
+    )
+
+
+def get_cognite_instance_ids(
+    instance_ids: list[TViewInstance],
+) -> list[dict[str, str]]:
+    return [
+        get_cognite_instance_id(instance_id) for instance_id in instance_ids
+    ]
+
+
+def get_cognite_instance_id(instance_id: TViewInstance) -> dict[str, str]:
+    return {"space": instance_id.space, "externalId": instance_id.external_id}

industrial_model/cognite_adapters/view_mapper.py

@@ -0,0 +1,16 @@
+from cognite.client.data_classes.data_modeling import (
+    View,
+)
+
+
+class ViewMapper:
+    def __init__(self, views: list[View]):
+        self._views_as_dict = {view.external_id: view for view in views}
+
+    def get_view(self, view_external_id: str) -> View:
+        if view_external_id not in self._views_as_dict:
+            raise ValueError(
+                f"View {view_external_id} is not available in data model"
+            )
+
+        return self._views_as_dict[view_external_id]

industrial_model/config.py

@@ -0,0 +1,10 @@
+from industrial_model.models import RootModel
+
+
+class DataModelId(RootModel):
+    external_id: str
+    space: str
+    version: str
+
+    def as_tuple(self) -> tuple[str, str, str]:
+        return self.space, self.external_id, self.version

industrial_model/constants.py

@@ -0,0 +1,24 @@
+from typing import Literal
+
+SORT_DIRECTION = Literal["ascending", "descending"]
+LEAF_EXPRESSION_OPERATORS = Literal[
+    "==",
+    "in",
+    ">=",
+    ">",
+    "<=",
+    "<",
+    "nested",
+    "exists",
+    "prefix",
+    "containsAll",
+    "containsAny",
+]
+BOOL_EXPRESSION_OPERATORS = Literal["not", "and", "or"]
+
+
+NESTED_SEP = "|"
+EDGE_MARKER = "<EdgeMarker>"
+EDGE_DIRECTION = Literal["outwards", "inwards"]
+MAX_LIMIT = 10_000
+DEFAULT_LIMIT = 5_000

industrial_model/engines/__init__.py

@@ -0,0 +1,4 @@
+from .async_engine import AsyncEngine
+from .engine import Engine
+
+__all__ = ["Engine", "AsyncEngine"]

industrial_model/engines/async_engine.py

@@ -0,0 +1,37 @@
+from cognite.client import CogniteClient
+
+from industrial_model.config import DataModelId
+from industrial_model.models import (
+    PaginatedResult,
+    TViewInstance,
+    ValidationMode,
+)
+from industrial_model.statements import Statement
+from industrial_model.utils import run_async
+
+from .engine import Engine
+
+
+class AsyncEngine:
+    def __init__(
+        self,
+        cognite_client: CogniteClient,
+        data_model_id: DataModelId,
+    ):
+        self._engine = Engine(cognite_client, data_model_id)
+
+    async def query_async(
+        self,
+        statement: Statement[TViewInstance],
+        validation_mode: ValidationMode = "raiseOnError",
+    ) -> PaginatedResult[TViewInstance]:
+        return await run_async(self._engine.query, statement, validation_mode)
+
+    async def query_all_pages_async(
+        self,
+        statement: Statement[TViewInstance],
+        validation_mode: ValidationMode = "raiseOnError",
+    ) -> list[TViewInstance]:
+        return await run_async(
+            self._engine.query_all_pages, statement, validation_mode
+        )

industrial_model/engines/engine.py

@@ -0,0 +1,62 @@
+from typing import Any
+
+from cognite.client import CogniteClient
+
+from industrial_model.cognite_adapters import CogniteAdapter
+from industrial_model.config import DataModelId
+from industrial_model.models import (
+    PaginatedResult,
+    TViewInstance,
+    ValidationMode,
+)
+from industrial_model.statements import Statement
+
+
+class Engine:
+    def __init__(
+        self,
+        cognite_client: CogniteClient,
+        data_model_id: DataModelId,
+    ):
+        self._cognite_adapter = CogniteAdapter(cognite_client, data_model_id)
+
+    def query(
+        self,
+        statement: Statement[TViewInstance],
+        validation_mode: ValidationMode = "raiseOnError",
+    ) -> PaginatedResult[TViewInstance]:
+        data, next_cursor = self._cognite_adapter.query(statement, False)
+
+        return PaginatedResult(
+            data=self._validate_data(statement.entity, data, validation_mode),
+            next_cursor=next_cursor,
+            has_next_page=next_cursor is not None,
+        )
+
+    def query_all_pages(
+        self,
+        statement: Statement[TViewInstance],
+        validation_mode: ValidationMode = "raiseOnError",
+    ) -> list[TViewInstance]:
+        if statement.cursor_:
+            raise ValueError("Cursor should be none when querying all pages")
+
+        data, _ = self._cognite_adapter.query(statement, True)
+
+        return self._validate_data(statement.entity, data, validation_mode)
+
+    def _validate_data(
+        self,
+        entity: type[TViewInstance],
+        data: list[dict[str, Any]],
+        validation_mode: ValidationMode,
+    ) -> list[TViewInstance]:
+        result: list[TViewInstance] = []
+        for item in data:
+            try:
+                result.append(entity.model_validate(item))
+            except Exception:
+                if validation_mode == "ignoreOnError":
+                    continue
+                raise
+        return result

industrial_model/models/__init__.py

@@ -0,0 +1,19 @@
+from .base import RootModel
+from .entities import (
+    InstanceId,
+    PaginatedResult,
+    TViewInstance,
+    ValidationMode,
+    ViewInstance,
+    ViewInstanceConfig,
+)
+
+__all__ = [
+    "RootModel",
+    "InstanceId",
+    "TViewInstance",
+    "ViewInstance",
+    "ValidationMode",
+    "PaginatedResult",
+    "ViewInstanceConfig",
+]

industrial_model/models/base.py

@@ -0,0 +1,46 @@
+from typing import (
+    Any,
+    dataclass_transform,
+)
+
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic._internal import _model_construction
+from pydantic.alias_generators import to_camel
+
+from industrial_model.statements import Column
+
+
+@dataclass_transform(kw_only_default=True, field_specifiers=(Field,))
+class DBModelMetaclass(_model_construction.ModelMetaclass):
+    is_constructing: bool = False
+
+    def __new__(
+        mcs,
+        cls_name: str,
+        bases: tuple[type[Any], ...],
+        namespace: dict[str, Any],
+        **kwargs: Any,
+    ) -> type:
+        mcs.is_constructing = True
+        cls = super().__new__(mcs, cls_name, bases, namespace, **kwargs)
+        mcs.is_constructing = False
+        return cls
+
+    def __getattr__(self, key: str) -> Any:
+        if self.is_constructing:
+            return super().__getattr__(key)  # type: ignore
+
+        try:
+            return super().__getattr__(key)  # type: ignore
+        except AttributeError:
+            if key in self.model_fields:
+                return Column(self.model_fields[key].alias or key)
+            raise
+
+
+class RootModel(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )

industrial_model/models/entities.py

@@ -0,0 +1,55 @@
+from typing import (
+    Any,
+    ClassVar,
+    Generic,
+    Literal,
+    TypedDict,
+    TypeVar,
+)
+
+from .base import DBModelMetaclass, RootModel
+
+
+class InstanceId(RootModel):
+    external_id: str
+    space: str
+
+    def __hash__(self) -> int:
+        return hash((self.external_id, self.space))
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            other is not None
+            and isinstance(other, InstanceId)
+            and self.external_id == other.external_id
+            and self.space == other.space
+        )
+
+    def as_tuple(self) -> tuple[str, str]:
+        return (self.space, self.external_id)
+
+
+class ViewInstanceConfig(TypedDict, total=False):
+    view_external_id: str | None
+    instance_spaces: list[str] | None
+    instance_spaces_prefix: str | None
+
+
+class ViewInstance(InstanceId, metaclass=DBModelMetaclass):
+    view_config: ClassVar[ViewInstanceConfig] = ViewInstanceConfig()
+
+    @classmethod
+    def get_view_external_id(cls) -> str:
+        return cls.view_config.get("view_external_id") or cls.__name__
+
+
+TViewInstance = TypeVar("TViewInstance", bound=ViewInstance)
+
+
+class PaginatedResult(RootModel, Generic[TViewInstance]):
+    data: list[TViewInstance]
+    has_next_page: bool
+    next_cursor: str | None
+
+
+ValidationMode = Literal["raiseOnError", "ignoreOnError"]

industrial_model/queries/__init__.py

@@ -0,0 +1,10 @@
+from .models import BasePaginatedQuery, BaseQuery
+from .params import NestedQueryParam, QueryParam, SortParam
+
+__all__ = [
+    "BaseQuery",
+    "BasePaginatedQuery",
+    "SortParam",
+    "QueryParam",
+    "NestedQueryParam",
+]

industrial_model/queries/models.py

@@ -0,0 +1,37 @@
+from industrial_model.models import RootModel, TViewInstance
+from industrial_model.statements import Statement
+
+from .params import NestedQueryParam, QueryParam, SortParam
+
+
+class BaseQuery(RootModel):
+    def to_statement(
+        self, entity: type[TViewInstance]
+    ) -> Statement[TViewInstance]:
+        statement = Statement(entity)
+
+        for key, item in self.__class__.model_fields.items():
+            values = getattr(self, key)
+            if not values:
+                continue
+            for metadata_item in item.metadata:
+                if isinstance(metadata_item, SortParam):
+                    statement.sort(values, metadata_item.direction)
+                elif isinstance(metadata_item, QueryParam | NestedQueryParam):
+                    statement.where(metadata_item.to_expression(values))
+
+        return statement
+
+
+class BasePaginatedQuery(BaseQuery):
+    limit: int = 1000
+    cursor: str | None = None
+
+    def to_statement(
+        self, entity: type[TViewInstance]
+    ) -> Statement[TViewInstance]:
+        statement = super().to_statement(entity)
+        statement.limit(self.limit)
+        statement.cursor(self.cursor)
+
+        return statement

industrial_model/queries/params.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from typing import Any
+
+from industrial_model.constants import (
+    LEAF_EXPRESSION_OPERATORS,
+    SORT_DIRECTION,
+)
+from industrial_model.statements import LeafExpression
+
+
+@dataclass
+class QueryParam:
+    property: str
+    operator: LEAF_EXPRESSION_OPERATORS
+
+    def to_expression(self, value: Any) -> LeafExpression:
+        if self.operator == "nested":
+            raise ValueError("Can not have nested operator on QuertParam")
+
+        return LeafExpression(
+            property=self.property,
+            operator=self.operator,
+            value=value,
+        )
+
+
+@dataclass
+class NestedQueryParam:
+    property: str
+    value: QueryParam
+
+    def to_expression(self, value: Any) -> LeafExpression:
+        return LeafExpression(
+            property=self.property,
+            operator="nested",
+            value=self.value.to_expression(value),
+        )
+
+
+@dataclass
+class SortParam:
+    direction: SORT_DIRECTION

industrial_model/statements/__init__.py

@@ -0,0 +1,70 @@
+from dataclasses import dataclass, field
+from typing import Any, Generic, Self, TypeVar
+
+from industrial_model.constants import DEFAULT_LIMIT, SORT_DIRECTION
+
+from .expressions import (
+    BoolExpression,
+    Column,
+    Expression,
+    LeafExpression,
+    and_,
+    col,
+    or_,
+)
+
+T = TypeVar("T")
+
+
+@dataclass
+class Statement(Generic[T]):
+    entity: type[T] = field(init=True)
+    where_clauses: list[Expression] = field(init=False, default_factory=list)
+    sort_clauses: list[tuple[Column, SORT_DIRECTION]] = field(
+        init=False, default_factory=list
+    )
+    limit_: int = field(init=False, default=DEFAULT_LIMIT)
+    cursor_: str | None = field(init=False, default=None)
+
+    def where(self, *expressions: bool | Expression) -> Self:
+        for expression in expressions:
+            assert isinstance(expression, Expression)
+            self.where_clauses.append(expression)
+        return self
+
+    def asc(self, property: Any) -> Self:
+        self.sort_clauses.append((Column(property), "ascending"))
+        return self
+
+    def desc(self, property: Any) -> Self:
+        self.sort_clauses.append((Column(property), "descending"))
+        return self
+
+    def sort(self, property: Any, direction: SORT_DIRECTION) -> Self:
+        self.sort_clauses.append((Column(property), direction))
+        return self
+
+    def limit(self, limit: int) -> Self:
+        self.limit_ = limit
+        return self
+
+    def cursor(self, cursor: str | None) -> Self:
+        self.cursor_ = cursor
+        return self
+
+
+def select(entity: type[T]) -> Statement[T]:
+    return Statement(entity)
+
+
+__all__ = [
+    "Statement",
+    "select",
+    "Column",
+    "col",
+    "Expression",
+    "LeafExpression",
+    "BoolExpression",
+    "and_",
+    "or_",
+]