orchestrator-core 3.1.2rc4__py3-none-any.whl → 3.2.0__py3-none-any.whl

orchestrator/domain/base.py

@@ -36,7 +36,7 @@ from more_itertools import bucket, first, flatten, one, only
 from pydantic import BaseModel, ConfigDict, Field, ValidationError
 from pydantic.fields import PrivateAttr
 from sqlalchemy import select
-from sqlalchemy.orm import selectinload
+from sqlalchemy.orm import joinedload, selectinload
 
 from orchestrator.db import (
     ProductBlockTable,
@@ -47,13 +47,19 @@ from orchestrator.db import (
     SubscriptionTable,
     db,
 )
-from orchestrator.domain.helpers import _to_product_block_field_type_iterable
+from orchestrator.db.queries.subscription_instance import get_subscription_instance_dict
+from orchestrator.domain.helpers import (
+    _to_product_block_field_type_iterable,
+    get_root_blocks_to_instance_ids,
+    no_private_attrs,
+)
 from orchestrator.domain.lifecycle import (
     ProductLifecycle,
     lookup_specialized_type,
     register_specialized_type,
     validate_lifecycle_status,
 )
+from orchestrator.domain.subscription_instance_transform import field_transformation_rules, transform_instance_fields
 from orchestrator.services.products import get_product_by_id
 from orchestrator.types import (
     SAFE_USED_BY_TRANSITIONS_FOR_STATUS,
@@ -98,6 +104,12 @@ class DomainModel(BaseModel):
     def __init_subclass__(cls, *args: Any, lifecycle: list[SubscriptionLifecycle] | None = None, **kwargs: Any) -> None:
         pass
 
+    def __eq__(self, other: Any) -> bool:
+        # PrivateAttr fields are excluded from both objects during the equality check.
+        # Added for #652 primarily because ProductBlockModel._db_model is now lazy loaded.
+        with no_private_attrs(self), no_private_attrs(other):
+            return super().__eq__(other)
+
     @classmethod
     def __pydantic_init_subclass__(
         cls,
@@ -310,12 +322,7 @@ class DomainModel(BaseModel):
         for product_block_field_name, product_block_field_type in cls._product_block_fields_.items():
             filter_func = match_domain_model_attr_if_possible(product_block_field_name)
 
-            product_block_model: Any = product_block_field_type
-            if is_list_type(product_block_field_type):
-                _origin, args = get_origin_and_args(product_block_field_type)
-                product_block_model = one(args)
-
-            possible_product_block_types = get_possible_product_block_types(product_block_model)
+            possible_product_block_types = flatten_product_block_types(product_block_field_type)
             field_type_names = list(possible_product_block_types.keys())
             filtered_instances = flatten([grouped_instances.get(name, []) for name in field_type_names])
             instance_list = list(filter(filter_func, filtered_instances))
@@ -458,6 +465,15 @@ class DomainModel(BaseModel):
             raise ValueError(f"Cannot link the same subscription instance multiple times: {details}")
 
 
+def flatten_product_block_types(product_block_field_type: Any) -> dict[str, type["ProductBlockModel"]]:
+    """Extract product block types and return mapping of product block names to product block classes."""
+    product_block_model: Any = product_block_field_type
+    if is_list_type(product_block_field_type):
+        _origin, args = get_origin_and_args(product_block_field_type)
+        product_block_model = one(args)
+    return get_possible_product_block_types(product_block_model)
+
+
 def get_depends_on_product_block_type_list(
     product_block_types: dict[str, type["ProductBlockModel"] | tuple[type["ProductBlockModel"]]],
 ) -> list[type["ProductBlockModel"]]:
@@ -522,7 +538,7 @@ class ProductBlockModel(DomainModel):
     product_block_id: ClassVar[UUID]
     description: ClassVar[str]
     tag: ClassVar[str]
-    _db_model: SubscriptionInstanceTable = PrivateAttr()
+    _db_model: SubscriptionInstanceTable | None = PrivateAttr(default=None)
 
     # Product block name. This needs to be an instance var because its part of the API (we expose it to the frontend)
     # Is actually optional since abstract classes don't have it.
@@ -681,7 +697,7 @@ class ProductBlockModel(DomainModel):
             **sub_instances,
             **kwargs,
         )
-        model._db_model = db_model
+        model.db_model = db_model
         return model
 
     @classmethod
@@ -737,7 +753,7 @@ class ProductBlockModel(DomainModel):
 
         cls._fix_pb_data()
         model = cls(**data)
-        model._db_model = other._db_model
+        model.db_model = other.db_model
         return model
 
     @classmethod
@@ -795,7 +811,7 @@ class ProductBlockModel(DomainModel):
                 **instance_values,  # type: ignore
                 **sub_instances,
             )
-            model._db_model = subscription_instance
+            model.db_model = subscription_instance
 
             return model
         except ValidationError:
@@ -913,14 +929,15 @@ class ProductBlockModel(DomainModel):
 
             # If this is a "foreign" instance we just stop saving and return it so only its relation is saved
             # We should not touch these themselves
-            if self.subscription and subscription_instance.subscription_id != subscription_id:
+            if self.owner_subscription_id != subscription_id:
                 return [], subscription_instance
 
-            self._db_model = subscription_instance
-        else:
-            subscription_instance = self._db_model
+            self.db_model = subscription_instance
+        elif subscription_instance := self.db_model:
             # We only need to add to the session if the subscription_instance does not exist.
             db.session.add(subscription_instance)
+        else:
+            raise ValueError("Cannot save ProductBlockModel without a db_model")
 
         subscription_instance.subscription_id = subscription_id
 
@@ -947,22 +964,32 @@ class ProductBlockModel(DomainModel):
         return sub_instances + [subscription_instance], subscription_instance
 
     @property
-    def subscription(self) -> SubscriptionTable:
-        return self.db_model.subscription
+    def subscription(self) -> SubscriptionTable | None:
+        return self.db_model.subscription if self.db_model else None
 
     @property
-    def db_model(self) -> SubscriptionInstanceTable:
+    def db_model(self) -> SubscriptionInstanceTable | None:
+        if not self._db_model:
+            self._db_model = db.session.execute(
+                select(SubscriptionInstanceTable).where(
+                    SubscriptionInstanceTable.subscription_instance_id == self.subscription_instance_id
+                )
+            ).scalar_one_or_none()
         return self._db_model
 
+    @db_model.setter
+    def db_model(self, value: SubscriptionInstanceTable) -> None:
+        self._db_model = value
+
     @property
-    def in_use_by(self) -> list[SubscriptionInstanceTable]:
+    def in_use_by(self) -> list[SubscriptionInstanceTable]:  # TODO check where used, might need eagerloading
         """This provides a list of product blocks that depend on this product block."""
-        return self._db_model.in_use_by
+        return self.db_model.in_use_by if self.db_model else []
 
     @property
-    def depends_on(self) -> list[SubscriptionInstanceTable]:
+    def depends_on(self) -> list[SubscriptionInstanceTable]:  # TODO check where used, might need eagerloading
         """This provides a list of product blocks that this product block depends on."""
-        return self._db_model.depends_on
+        return self.db_model.depends_on if self.db_model else []
 
 
 class ProductModel(BaseModel):
@@ -1010,9 +1037,11 @@ class SubscriptionModel(DomainModel):
         >>> SubscriptionInactive.from_subscription(subscription_id)  # doctest:+SKIP
     """
 
+    __model_dump_cache__: ClassVar[dict[UUID, "SubscriptionModel"] | None] = None
+
     product: ProductModel
     customer_id: str
-    _db_model: SubscriptionTable = PrivateAttr()
+    _db_model: SubscriptionTable | None = PrivateAttr(default=None)
     subscription_id: UUID = Field(default_factory=uuid4)  # pragma: no mutate
     description: str = "Initial subscription"  # pragma: no mutate
     status: SubscriptionLifecycle = SubscriptionLifecycle.INITIAL  # pragma: no mutate
@@ -1103,6 +1132,63 @@ class SubscriptionModel(DomainModel):
 
         return missing_data
 
+    @classmethod
+    def _load_root_instances(
+        cls,
+        subscription_id: UUID | UUIDstr,
+    ) -> dict[str, Optional[dict] | list[dict]]:
+        """Load root subscription instance(s) for this subscription model.
+
+        When a new subscription model is loaded from an existing subscription, this function loads the entire root
+        subscription instance(s) from database using an optimized postgres function. The result of that function
+        is used to instantiate the root product block(s).
+
+        The "old" method DomainModel._load_instances() would recursively load subscription instances from the
+        database and individually instantiate nested blocks, more or less "manually" reconstructing the subscription.
+
+        The "new" method SubscriptionModel._load_root_instances() takes a different approach; since it has all
+        data for the root subscription instance, it can rely on Pydantic to instantiate the root block and all
+        nested blocks in one go. This is also why it does not have the params `status` and `match_domain_attr` because
+        this information is already encoded in the domain model of a product.
+        """
+        root_block_instance_ids = get_root_blocks_to_instance_ids(subscription_id)
+
+        root_block_types = {
+            field_name: list(flatten_product_block_types(product_block_type).keys())
+            for field_name, product_block_type in cls._product_block_fields_.items()
+        }
+
+        def get_instances_by_block_names(block_names: list[str]) -> Iterable[dict]:
+            for block_name in block_names:
+                for instance_id in root_block_instance_ids.get(block_name, []):
+                    yield get_subscription_instance_dict(instance_id)
+
+        # Map root product block fields to subscription instance(s) dicts
+        instances = {
+            field_name: list(get_instances_by_block_names(block_names))
+            for field_name, block_names in root_block_types.items()
+        }
+
+        # Transform values according to domain models (list[dict] -> dict, add None as default for optionals)
+        rules = {
+            klass.name: field_transformation_rules(klass) for klass in ProductBlockModel.registry.values() if klass.name
+        }
+        for instance_list in instances.values():
+            for instance in instance_list:
+                transform_instance_fields(rules, instance)
+
+        # Support the (theoretical?) usecase of a list of root product blocks
+        def unpack_instance_list(field_name: str, instance_list: list[dict]) -> list[dict] | dict | None:
+            field_type = cls._product_block_fields_[field_name]
+            if is_list_type(field_type):
+                return instance_list
+            return only(instance_list)
+
+        return {
+            field_name: unpack_instance_list(field_name, instance_list)
+            for field_name, instance_list in instances.items()
+        }
+
     @classmethod
     def from_product_id(
         cls: type[S],
@@ -1168,7 +1254,7 @@ class SubscriptionModel(DomainModel):
             **fixed_inputs,
             **instances,
         )
-        model._db_model = subscription
+        model.db_model = subscription
         return model
 
     @classmethod
@@ -1201,17 +1287,26 @@ class SubscriptionModel(DomainModel):
             data["end_date"] = nowtz()
 
         model = cls(**data)
-        model._db_model = other._db_model
+        model.db_model = other._db_model
 
         return model
 
     # Some common functions shared by from_other_product and from_subscription
     @classmethod
-    def _get_subscription(cls: type[S], subscription_id: UUID | UUIDstr) -> Any:
-        return db.session.get(
-            SubscriptionTable,
-            subscription_id,
-            options=[
+    def _get_subscription(cls: type[S], subscription_id: UUID | UUIDstr) -> SubscriptionTable | None:
+        from orchestrator.settings import app_settings
+
+        if not isinstance(subscription_id, UUID | UUIDstr):
+            raise TypeError(f"subscription_id is of type {type(subscription_id)} instead of UUID | UUIDstr")
+
+        if app_settings.ENABLE_SUBSCRIPTION_MODEL_OPTIMIZATIONS:
+            # TODO #900 remove toggle and make this path the default
+            loaders = [
+                joinedload(SubscriptionTable.product).selectinload(ProductTable.fixed_inputs),
+            ]
+
+        else:
+            loaders = [
                 selectinload(SubscriptionTable.instances)
                 .joinedload(SubscriptionInstanceTable.product_block)
                 .selectinload(ProductBlockTable.resource_types),
@@ -1219,8 +1314,9 @@ class SubscriptionModel(DomainModel):
                     SubscriptionInstanceTable.in_use_by_block_relations
                 ),
                 selectinload(SubscriptionTable.instances).selectinload(SubscriptionInstanceTable.values),
-            ],
-        )
+            ]
+
+        return db.session.get(SubscriptionTable, subscription_id, options=loaders)
 
     @classmethod
     def _to_product_model(cls: type[S], product: ProductTable) -> ProductModel:
@@ -1246,7 +1342,9 @@ class SubscriptionModel(DomainModel):
         if not db_product:
             raise KeyError("Could not find a product for the given product_id")
 
-        subscription = cls._get_subscription(old_instantiation.subscription_id)
+        old_subscription_id = old_instantiation.subscription_id
+        if not (subscription := cls._get_subscription(old_subscription_id)):
+            raise ValueError(f"Subscription with id: {old_subscription_id}, does not exist")
         product = cls._to_product_model(db_product)
 
         status = SubscriptionLifecycle(subscription.status)
@@ -1266,6 +1364,7 @@ class SubscriptionModel(DomainModel):
             name, product_block = new_root
             instances = {name: product_block}
         else:
+            # TODO test using cls._load_root_instances() here as well
             instances = cls._load_instances(subscription.instances, status, match_domain_attr=False)  # type:ignore
 
         try:
@@ -1283,7 +1382,7 @@ class SubscriptionModel(DomainModel):
                 **fixed_inputs,
                 **instances,
             )
-            model._db_model = subscription
+            model.db_model = subscription
             return model
         except ValidationError:
             logger.exception(
@@ -1294,8 +1393,13 @@ class SubscriptionModel(DomainModel):
     @classmethod
     def from_subscription(cls: type[S], subscription_id: UUID | UUIDstr) -> S:
         """Use a subscription_id to return required fields of an existing subscription."""
-        subscription = cls._get_subscription(subscription_id)
-        if subscription is None:
+        from orchestrator.domain.context_cache import get_from_cache, store_in_cache
+        from orchestrator.settings import app_settings
+
+        if cached_model := get_from_cache(subscription_id):
+            return cast(S, cached_model)
+
+        if not (subscription := cls._get_subscription(subscription_id)):
             raise ValueError(f"Subscription with id: {subscription_id}, does not exist")
         product = cls._to_product_model(subscription.product)
 
@@ -1317,7 +1421,12 @@ class SubscriptionModel(DomainModel):
 
         fixed_inputs = {fi.name: fi.value for fi in subscription.product.fixed_inputs}
 
-        instances = cls._load_instances(subscription.instances, status, match_domain_attr=False)
+        instances: dict[str, Any]
+        if app_settings.ENABLE_SUBSCRIPTION_MODEL_OPTIMIZATIONS:
+            # TODO #900 remove toggle and make this path the default
+            instances = cls._load_root_instances(subscription_id)
+        else:
+            instances = cls._load_instances(subscription.instances, status, match_domain_attr=False)
 
         try:
             model = cls(
@@ -1334,7 +1443,10 @@ class SubscriptionModel(DomainModel):
                 **fixed_inputs,
                 **instances,
             )
-            model._db_model = subscription
+            model.db_model = subscription
+
+            store_in_cache(model)
+
             return model
         except ValidationError:
             logger.exception(
@@ -1350,7 +1462,7 @@ class SubscriptionModel(DomainModel):
                 f"Lifecycle status {self.status.value} requires specialized type {specialized_type!r}, was: {type(self)!r}"
             )
 
-        sub = db.session.get(
+        existing_sub = db.session.get(
             SubscriptionTable,
             self.subscription_id,
             options=[
@@ -1360,13 +1472,13 @@ class SubscriptionModel(DomainModel):
                 selectinload(SubscriptionTable.instances).selectinload(SubscriptionInstanceTable.values),
             ],
         )
-        if not sub:
-            sub = self._db_model
+        if not (sub := (existing_sub or self.db_model)):
+            raise ValueError("Cannot save SubscriptionModel without a db_model")
 
         # Make sure we refresh the object and not use an already mapped object
         db.session.refresh(sub)
 
-        self._db_model = sub
+        self.db_model = sub
         sub.product_id = self.product.product_id
         sub.customer_id = self.customer_id
         sub.description = self.description
@@ -1404,9 +1516,15 @@ class SubscriptionModel(DomainModel):
         db.session.flush()
 
     @property
-    def db_model(self) -> SubscriptionTable:
+    def db_model(self) -> SubscriptionTable | None:
+        if not self._db_model:
+            self._db_model = self._get_subscription(self.subscription_id)
         return self._db_model
 
+    @db_model.setter
+    def db_model(self, value: SubscriptionTable) -> None:
+        self._db_model = value
+
 
 def validate_base_model(
     name: str, cls: type[Any], base_model: type[BaseModel] = DomainModel, errors: list[str] | None = None

orchestrator/domain/context_cache.py

@@ -0,0 +1,62 @@
+# Copyright 2019-2025 SURF.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import contextlib
+from contextvars import ContextVar
+from typing import Iterator
+from uuid import UUID
+
+from orchestrator.domain import SubscriptionModel
+from pydantic_forms.types import UUIDstr
+
+__subscription_model_cache: ContextVar[dict[UUID, SubscriptionModel] | None] = ContextVar(
+    "subscription_model_cache", default=None
+)
+
+
+@contextlib.contextmanager
+def cache_subscription_models() -> Iterator:
+    """Caches SubscriptionModels for the duration of the context.
+
+    Inside this context, calling SubscriptionModel.from_subscription() twice with the same
+    subscription id will return the same instance.
+
+    The primary usecase is to improve performance of `@computed_field` properties on product blocks
+    which load other subscriptions.
+
+    Example usage:
+        subscription = SubscriptionModel.from_subscription("...")
+        with cache_subscription_models():
+           subscription_dict = subscription.model_dump()
+    """
+    before = __subscription_model_cache.set({})
+    try:
+        yield
+    finally:
+        __subscription_model_cache.reset(before)
+
+
+def get_from_cache(subscription_id: UUID | UUIDstr) -> SubscriptionModel | None:
+    """Retrieve SubscriptionModel from cache, if present."""
+    if (cache := __subscription_model_cache.get()) is None:
+        return None
+
+    id_ = subscription_id if isinstance(subscription_id, UUID) else UUID(subscription_id)
+    return cache.get(id_, None)
+
+
+def store_in_cache(model: SubscriptionModel) -> None:
+    """Store SubscriptionModel in cache, if required."""
+    if (cache := __subscription_model_cache.get()) is None:
+        return
+
+    cache[model.subscription_id] = model

orchestrator/domain/helpers.py

@@ -1,6 +1,15 @@
-from collections.abc import Iterable
+import contextlib
+from collections.abc import Iterable, Iterator
+from typing import Any
+from uuid import UUID
 
+from pydantic import BaseModel
+from sqlalchemy import select
+
+from orchestrator.db import ProductBlockTable, SubscriptionInstanceTable, db
 from orchestrator.types import filter_nonetype, get_origin_and_args, is_union_type
+from orchestrator.utils.functional import group_by_key
+from pydantic_forms.types import UUIDstr
 
 
 def _to_product_block_field_type_iterable(product_block_field_type: type | tuple[type]) -> Iterable[type]:
@@ -21,3 +30,34 @@ def _to_product_block_field_type_iterable(product_block_field_type: type | tuple
         return product_block_field_type
 
     return [product_block_field_type]
+
+
+@contextlib.contextmanager
+def no_private_attrs(model: Any) -> Iterator:
+    """PrivateAttrs from the given pydantic BaseModel are removed for the duration of this context."""
+    if not isinstance(model, BaseModel):
+        yield
+        return
+    private_attrs_reference = model.__pydantic_private__
+    try:
+        model.__pydantic_private__ = {}
+        yield
+    finally:
+        model.__pydantic_private__ = private_attrs_reference
+
+
+def get_root_blocks_to_instance_ids(subscription_id: UUID | UUIDstr) -> dict[str, list[UUID]]:
+    """Returns mapping of root product block names to list of subscription instance ids.
+
+    While recommended practice is to have only 1 root product block, it is possible to have multiple blocks or even a
+    list of root blocks. This function supports that.
+    """
+    block_name_to_instance_id_rows = db.session.execute(
+        select(ProductBlockTable.name, SubscriptionInstanceTable.subscription_instance_id)
+        .select_from(SubscriptionInstanceTable)
+        .join(ProductBlockTable)
+        .where(SubscriptionInstanceTable.subscription_id == subscription_id)
+        .order_by(ProductBlockTable.name)
+    ).all()
+
+    return group_by_key(block_name_to_instance_id_rows)  # type: ignore[arg-type]

orchestrator/domain/subscription_instance_transform.py

@@ -0,0 +1,114 @@
+# Copyright 2019-2025 SURF.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Functions to transform result of query SubscriptionInstanceAsJsonFunction to match the ProductBlockModel."""
+
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Iterable
+
+from more_itertools import first, only
+
+from orchestrator.types import is_list_type, is_optional_type
+
+if TYPE_CHECKING:
+    from orchestrator.domain.base import ProductBlockModel
+
+
+def _ensure_list(instance_or_value_list: Any) -> Any:
+    if instance_or_value_list is None:
+        return []
+
+    return instance_or_value_list
+
+
+def _instance_list_to_dict(product_block_field_type: type, instance_list: Any) -> Any:
+    if instance_list is None:
+        return None
+
+    match instance_list:
+        case list():
+            if instance := only(instance_list):
+                return instance
+
+            if not is_optional_type(product_block_field_type):
+                raise ValueError("Required subscription instance is missing in database")
+
+            return None  # Set the optional product block field to None
+        case _:
+            raise ValueError(f"All subscription instances should be returned as list, found {type(instance_list)}")  #
+
+
+def _value_list_to_value(field_type: type, value_list: Any) -> Any:
+    if value_list is None:
+        return None
+
+    match value_list:
+        case list():
+            if (value := only(value_list)) is not None:
+                return value
+
+            if not is_optional_type(field_type):
+                raise ValueError("Required subscription value is missing in database")
+
+            return None  # Set the optional resource type field to None
+        case _:
+            raise ValueError(f"All instance values should be returned as list, found {type(value_list)}")
+
+
+def field_transformation_rules(klass: type["ProductBlockModel"]) -> dict[str, Callable]:
+    """Create mapping of transformation rules for the given product block type."""
+
+    def create_rules() -> Iterable[tuple[str, Callable]]:
+        for field_name, product_block_field_type in klass._product_block_fields_.items():
+            if is_list_type(product_block_field_type):
+                yield field_name, _ensure_list
+            else:
+                yield field_name, partial(_instance_list_to_dict, product_block_field_type)
+
+        for field_name, field_type in klass._non_product_block_fields_.items():
+            if is_list_type(field_type):
+                yield field_name, _ensure_list
+            else:
+                yield field_name, partial(_value_list_to_value, field_type)
+
+    return dict(create_rules())
+
+
+def transform_instance_fields(all_rules: dict[str, dict[str, Callable]], instance: dict) -> None:
+    """Apply transformation rules to the given subscription instance dict."""
+
+    from orchestrator.domain.base import ProductBlockModel
+
+    # Lookup applicable rules through product block name
+    field_rules = all_rules[instance["name"]]
+
+    klass = ProductBlockModel.registry[instance["name"]]
+
+    # Ensure the product block's metadata is loaded
+    klass._fix_pb_data()
+
+    # Transform all fields in this subscription instance
+    try:
+        for field_name, rewrite_func in field_rules.items():
+            field_value = instance.get(field_name)
+            instance[field_name] = rewrite_func(field_value)
+    except ValueError as e:
+        raise ValueError(f"Invalid subscription instance data {instance}") from e
+
+    # Recurse into nested subscription instances
+    for field_value in instance.values():
+        if isinstance(field_value, dict):
+            transform_instance_fields(all_rules, field_value)
+        if isinstance(field_value, list) and isinstance(first(field_value, None), dict):
+            for list_value in field_value:
+                transform_instance_fields(all_rules, list_value)

orchestrator/graphql/resolvers/process.py

@@ -34,7 +34,7 @@ from orchestrator.graphql.utils import (
     is_querying_page_data,
     to_graphql_result_page,
 )
-from orchestrator.graphql.utils.get_query_loaders import get_query_loaders
+from orchestrator.graphql.utils.get_query_loaders import get_query_loaders_for_gql_fields
 from orchestrator.schemas.process import ProcessSchema
 from orchestrator.services.processes import load_process
 from orchestrator.utils.enrich_process import enrich_process
@@ -56,7 +56,7 @@ def _enrich_process(process: ProcessTable, with_details: bool = False) -> Proces
 
 
 async def resolve_process(info: OrchestratorInfo, process_id: UUID) -> ProcessType | None:
-    query_loaders = get_query_loaders(info, ProcessTable)
+    query_loaders = get_query_loaders_for_gql_fields(ProcessTable, info)
     stmt = select(ProcessTable).options(*query_loaders).where(ProcessTable.process_id == process_id)
     if process := db.session.scalar(stmt):
         is_detailed = _is_process_detailed(info)
@@ -83,7 +83,7 @@ async def resolve_processes(
         .selectinload(ProcessSubscriptionTable.subscription)
         .joinedload(SubscriptionTable.product)
     ]
-    query_loaders = get_query_loaders(info, ProcessTable) or default_loaders
+    query_loaders = get_query_loaders_for_gql_fields(ProcessTable, info) or default_loaders
     select_stmt = select(ProcessTable).options(*query_loaders)
     select_stmt = filter_processes(select_stmt, pydantic_filter_by, _error_handler)
     if query is not None:

orchestrator/graphql/resolvers/product.py

@@ -13,7 +13,7 @@ from orchestrator.graphql.resolvers.helpers import rows_from_statement
 from orchestrator.graphql.schemas.product import ProductType
 from orchestrator.graphql.types import GraphqlFilter, GraphqlSort, OrchestratorInfo
 from orchestrator.graphql.utils import create_resolver_error_handler, is_querying_page_data, to_graphql_result_page
-from orchestrator.graphql.utils.get_query_loaders import get_query_loaders
+from orchestrator.graphql.utils.get_query_loaders import get_query_loaders_for_gql_fields
 from orchestrator.utils.search_query import create_sqlalchemy_select
 
 logger = structlog.get_logger(__name__)
@@ -33,7 +33,7 @@ async def resolve_products(
     pydantic_sort_by: list[Sort] = [item.to_pydantic() for item in sort_by] if sort_by else []
     logger.debug("resolve_products() called", range=[after, after + first], sort=sort_by, filter=pydantic_filter_by)
 
-    query_loaders = get_query_loaders(info, ProductTable)
+    query_loaders = get_query_loaders_for_gql_fields(ProductTable, info)
     select_stmt = select(ProductTable).options(*query_loaders)
     select_stmt = filter_products(select_stmt, pydantic_filter_by, _error_handler)