Rework SubscriptionModel.from_subscription and in_use_by_id injection to remove dependency on Redis (#880)

* v1: ~45ms query (incomplete)

- pgsql:
  - blocks: return all as list, optionals not included
  - values: return only scalars (incomplete), optionals not included
- base.py:
  - blocks: rewrite single block relations from list to dict, add default list for missing optionals
  - values: add default list/none for missing optionals

* v2: ~100ms query

- pgsql:
  - blocks: return all as list, optionals not included
  - values: return as scalars or list depending on mapping, add default
    array/null for missing optionals
- base.py:
  - blocks: rewrite single block relations from list to dict, add default list for missing optionals
  - values: no rewrite needed

* v3: ~48ms query

- pgsql:
  - blocks: return all as list, optionals not included
  - values: return all as list, optionals not included
- base.py:
  - blocks: rewrite single block relations from list to dict, add default list for missing optionals
  - values: rewrite single values from list to dict, add default
    list/None for missing optionals

* Refactor get_query_loaders function to be reusable

* v3.1

- Further improve SubscriptionModel.from_subscription
- UNION ALL slightly better than UNION when no duplicates are possible
- Temporary toggle ENABLE_SUBSCRIPTION_MODEL_OPTIMIZATIONS

* v3.2

- Speed up in_use_by injection in build_extended_domain_model

* Update tests

* Rename function and actually use the sqlalchemy definition

* Sort value/instance lists ascending to be consistent with sqlalchemy

* Also use ENABLE_SUBSCRIPTION_MODEL_OPTIMIZATIONS to toggle subscription model caching

* Add logic to SubscriptionModel._load_root_instances and document it, refactor to support list of root product blocks

* Lots of fixes + code reorganization

* Add headers

* Fix flaky test

* Add comment and fix type ignore

* Raise clear TypeError when passing an invalid subscription_id

* Add test test_subscription_detail_domain_model_optimizations for code coverage

* Address review comment

* Move serializable property tests to separate file

* Add testcase for retrieving SubscriptionModel in serializable property

* Add todo's for removing ENABLE_SUBSCRIPTION_MODEL_OPTIMIZATIONS

* Add context_cache.cache_subscription_models

* Document __eq__ override

* Ensure sqlalchemy model loaders are always initialized

* Bump version to 3.2.0rc1

Author: Mark90

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 3.1.2
+current_version = 3.2.0rc1
 commit = False
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(rc(?P<build>\d+))?

orchestrator/__init__.py

@@ -13,7 +13,7 @@
 
 """This is the orchestrator workflow engine."""
 
-__version__ = "3.1.2"
+__version__ = "3.2.0rc1"
 
 from orchestrator.app import OrchestratorCore
 from orchestrator.settings import app_settings

orchestrator/db/loaders.py

@@ -1,5 +1,5 @@
 from functools import reduce
-from typing import Any, Callable, Iterator, NamedTuple, cast
+from typing import Any, Callable, Iterable, Iterator, NamedTuple, cast
 
 import structlog
 from sqlalchemy import inspect
@@ -131,12 +131,15 @@ def init_model_loaders() -> None:
         _MODEL_LOADERS[model] = dict(_inspect_model(model))
 
 
-def lookup_attr_loaders(model: type[DbBaseModel], attr: str) -> list[AttrLoader]:
+def _lookup_attr_loaders(model: type[DbBaseModel], attr: str) -> list[AttrLoader]:
     """Return loader(s) for an attribute on the given model."""
+    if not _MODEL_LOADERS:
+        # Ensure loaders are always initialized
+        init_model_loaders()
     return _MODEL_LOADERS.get(model, {}).get(attr, [])
 
 
-def join_attr_loaders(loaders: list[AttrLoader]) -> Load | None:
+def _join_attr_loaders(loaders: list[AttrLoader]) -> Load | None:
     """Given 1 or more attribute loaders, instantiate and chain them together."""
     if not loaders:
         return None
@@ -151,3 +154,48 @@ def chain_loader_func(final_loader: Load, next: AttrLoader) -> Load:
         return getattr(final_loader, next.loader_fn.__name__)(next.attr)
 
     return reduce(chain_loader_func, other_loaders, loader_fn)
+
+
+def _split_path(query_path: str) -> Iterable[str]:
+    yield from (field for field in query_path.split("."))
+
+
+def get_query_loaders_for_model_paths(root_model: type[DbBaseModel], model_paths: list[str]) -> list[Load]:
+    """Get sqlalchemy query loaders to use for the model based on the paths."""
+    # Sort by length to find the longest match first
+    model_paths.sort(key=lambda x: x.count("."), reverse=True)
+
+    def get_loader_for_path(model_path: str) -> tuple[str, Load | None]:
+        next_model = root_model
+
+        matched_fields: list[str] = []
+        path_loaders: list[AttrLoader] = []
+
+        for field in _split_path(model_path):
+            if not (attr_loaders := _lookup_attr_loaders(next_model, field)):
+                break
+
+            matched_fields.append(field)
+            path_loaders.extend(attr_loaders)
+            next_model = attr_loaders[-1].next_model
+
+        return ".".join(matched_fields), _join_attr_loaders(path_loaders)
+
+    query_loaders: dict[str, Load] = {}
+
+    for path in model_paths:
+        matched_path, loader = get_loader_for_path(path)
+        if not matched_path or not loader or matched_path in query_loaders:
+            continue
+        if any(known_path.startswith(f"{matched_path}.") for known_path in query_loaders):
+            continue
+        query_loaders[matched_path] = loader
+
+    loaders = list(query_loaders.values())
+    logger.debug(
+        "Generated query loaders for paths",
+        root_model=root_model,
+        model_paths=model_paths,
+        query_loaders=[str(i.path) for i in loaders],
+    )
+    return loaders

orchestrator/db/models.py

@@ -15,6 +15,7 @@
 
 import enum
 from datetime import datetime, timezone
+from uuid import UUID
 
 import sqlalchemy
 import structlog
@@ -43,6 +44,7 @@
 from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.ext.orderinglist import ordering_list
 from sqlalchemy.orm import Mapped, deferred, mapped_column, object_session, relationship, undefer
+from sqlalchemy.sql.functions import GenericFunction
 from sqlalchemy_utils import TSVectorType, UUIDType
 
 from orchestrator.config.assignee import Assignee
@@ -667,3 +669,14 @@ class EngineSettingsTable(BaseModel):
     global_lock = mapped_column(Boolean(), default=False, nullable=False, primary_key=True)
     running_processes = mapped_column(Integer(), default=0, nullable=False)
     __table_args__: tuple = (CheckConstraint(running_processes >= 0, name="check_running_processes_positive"), {})
+
+
+class SubscriptionInstanceAsJsonFunction(GenericFunction):
+    # Added in migration 42b3d076a85b
+    name = "subscription_instance_as_json"
+
+    type = pg.JSONB()
+    inherit_cache = True
+
+    def __init__(self, sub_inst_id: UUID):
+        super().__init__(sub_inst_id)

orchestrator/db/queries/__init__.py

@@ -0,0 +1,85 @@
+# 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.
+
+from uuid import UUID
+
+from sqlalchemy import UUID as SA_UUID
+from sqlalchemy import cast as sa_cast
+from sqlalchemy import select
+from sqlalchemy.orm import raiseload
+
+from orchestrator.db import SubscriptionInstanceRelationTable, SubscriptionInstanceTable, db
+
+
+def _eagerload_subscription_instances(
+    subscription_id: UUID | str, instance_attributes: list[str]
+) -> list[SubscriptionInstanceTable]:
+    """Given a subscription id, recursively query all depends_on subscription instances with the instance_attributes eagerloaded.
+
+    Note: accessing instance attributes on the result that were not explicitly loaded will
+    trigger a sqlalchemy error.
+    """
+    from orchestrator.db.loaders import get_query_loaders_for_model_paths
+
+    # CTE to recursively get all subscription instance ids the subscription depends on
+    instance_ids_cte = (
+        select(
+            sa_cast(None, SA_UUID(as_uuid=True)).label("in_use_by_id"),
+            SubscriptionInstanceTable.subscription_instance_id.label("depends_on_id"),
+        )
+        .where(SubscriptionInstanceTable.subscription_id == subscription_id)
+        .cte(name="recursive_instance_ids", recursive=True)
+    )
+
+    cte_alias = instance_ids_cte.alias()
+    rel_alias = select(SubscriptionInstanceRelationTable).alias()
+
+    instance_ids = instance_ids_cte.union(
+        select(rel_alias.c.in_use_by_id, rel_alias.c.depends_on_id).where(
+            rel_alias.c.in_use_by_id == cte_alias.c.depends_on_id
+        )
+    )
+
+    select_all_instance_ids = select(instance_ids.c.depends_on_id).subquery()
+
+    # Eagerload specified instance attributes
+    query_loaders = get_query_loaders_for_model_paths(SubscriptionInstanceTable, instance_attributes)
+    # Prevent unwanted lazyloading of all other attributes
+    query_loaders += [raiseload("*")]  # type: ignore[list-item]  # todo fix this type
+    stmt = (
+        select(SubscriptionInstanceTable)
+        .where(SubscriptionInstanceTable.subscription_instance_id.in_(select(select_all_instance_ids)))
+        .options(*query_loaders)
+    )
+
+    return db.session.scalars(stmt).all()  # type: ignore[return-value]  # todo fix this type
+
+
+def eagerload_all_subscription_instances(subscription_id: UUID | str) -> list[SubscriptionInstanceTable]:
+    """Recursively find the subscription's depends_on instances and resolve relations for SubscriptionModel.from_subscription_id()."""
+    instance_attributes = [
+        "subscription.product",
+        "product_block",
+        "values.resource_type",
+        "depends_on",
+        "in_use_by",
+    ]
+    return _eagerload_subscription_instances(subscription_id, instance_attributes)
+
+
+def eagerload_all_subscription_instances_only_inuseby(subscription_id: UUID | str) -> list[SubscriptionInstanceTable]:
+    """Recursively find the subscription's depends_on instances and resolve their in_use_by relations."""
+    instance_attributes = [
+        "in_use_by",
+    ]
+    return _eagerload_subscription_instances(subscription_id, instance_attributes)

orchestrator/db/queries/subscription.py

@@ -0,0 +1,28 @@
+# 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.
+
+from uuid import UUID
+
+from sqlalchemy import select
+
+from orchestrator.db import db
+from orchestrator.db.models import SubscriptionInstanceAsJsonFunction
+
+
+def get_subscription_instance_dict(subscription_instance_id: UUID) -> dict:
+    """Query the subscription instance as aggregated JSONB and returns it as a dict.
+
+    Note: all values are returned as lists and have to be transformed by the caller.
+    It was attempted to do this in the DB query but this gave worse performance.
+    """
+    return db.session.execute(select(SubscriptionInstanceAsJsonFunction(subscription_instance_id))).scalar_one()