Reduce DB queries in GQL resolver for subscriptions (product, inuseby/dependson) (#745)

* Eagerload SubscriptionTable.product in initial query to avoid lazy loading

Signed-off-by: Mark90 <[email protected]>

* Add dataloaders for in_use_by and depends_on subscriptions

Signed-off-by: Mark90 <[email protected]>

* Add custom_context_getter option to register_graphql()

Signed-off-by: Mark90 <[email protected]>

* Bump version to 2.7.6rc1

---------

Signed-off-by: Mark90 <[email protected]>

Author: Mark90

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 2.7.5
+current_version = 2.7.6rc1
 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__ = "2.7.5"
+__version__ = "2.7.6rc1"
 
 from orchestrator.app import OrchestratorCore
 from orchestrator.settings import app_settings

orchestrator/app.py

@@ -50,6 +50,7 @@
 from orchestrator.domain import SUBSCRIPTION_MODEL_REGISTRY, SubscriptionModel
 from orchestrator.exception_handlers import problem_detail_handler
 from orchestrator.graphql import Mutation, Query, create_graphql_router
+from orchestrator.graphql.schema import ContextGetterFactory
 from orchestrator.graphql.schemas.subscription import SubscriptionInterface
 from orchestrator.graphql.types import ScalarOverrideType, StrawberryModelType
 from orchestrator.log_config import LOGGER_OVERRIDES
@@ -205,6 +206,7 @@ def register_graphql(
         graphql_models: StrawberryModelType | None = None,
         scalar_overrides: ScalarOverrideType | None = None,
         extensions: list | None = None,
+        custom_context_getter: ContextGetterFactory | None = None,
     ) -> None:
         new_router = create_graphql_router(
             self.auth_manager,
@@ -216,6 +218,7 @@ def register_graphql(
             graphql_models,
             scalar_overrides,
             extensions=extensions,
+            custom_context_getter=custom_context_getter,
         )
         if not self.graphql_router:
             self.graphql_router = new_router

orchestrator/graphql/__init__.py

@@ -24,7 +24,7 @@
     OrchestratorSchema,
     Query,
     create_graphql_router,
-    get_context,
+    default_context_getter,
 )
 from orchestrator.graphql.schemas import DEFAULT_GRAPHQL_MODELS
 from orchestrator.graphql.types import SCALAR_OVERRIDES
@@ -38,7 +38,7 @@
     "Mutation",
     "OrchestratorGraphqlRouter",
     "OrchestratorSchema",
-    "get_context",
+    "default_context_getter",
     "create_graphql_router",
     "EnumDict",
     "add_class_to_strawberry",

orchestrator/graphql/loaders/__init__.py

@@ -0,0 +1,224 @@
+from itertools import chain
+from typing import Any, NamedTuple
+from uuid import UUID
+
+import structlog
+from sqlalchemy import Row, select
+from sqlalchemy import Text as SaText
+from sqlalchemy import cast as sa_cast
+from sqlalchemy.orm import aliased
+from strawberry.dataloader import DataLoader
+
+from orchestrator.db import (
+    ResourceTypeTable,
+    SubscriptionInstanceTable,
+    SubscriptionInstanceValueTable,
+    SubscriptionTable,
+    db,
+)
+from orchestrator.db.models import (
+    SubscriptionInstanceRelationTable,
+)
+from orchestrator.services.subscriptions import RELATION_RESOURCE_TYPES
+from orchestrator.types import SubscriptionLifecycle
+
+logger = structlog.get_logger(__name__)
+
+
+class Relation(NamedTuple):
+    depends_on_sub_id: UUID
+    in_use_by_sub_id: UUID
+
+
+def _get_instance_relations(instance_relations_query: Any) -> list[Relation]:
+    def to_relation(row: Row[Any]) -> Relation:
+        return Relation(row[0], row[1])
+
+    return [to_relation(row) for row in db.session.execute(instance_relations_query)]
+
+
+async def _get_in_use_by_instance_relations(subscription_ids: list[UUID], filter_statuses: list[str]) -> list[Relation]:
+    """Get in_use_by by relations through subscription instance hierarchy."""
+    in_use_by_subscriptions = aliased(SubscriptionTable)
+    in_use_by_instances = aliased(SubscriptionInstanceTable)
+    depends_on_instances = aliased(SubscriptionInstanceTable)
+
+    query_get_in_use_by_ids = (
+        select(depends_on_instances.subscription_id, in_use_by_instances.subscription_id)
+        .distinct()
+        .join(in_use_by_instances.subscription)
+        .join(in_use_by_instances.depends_on_block_relations)
+        .join(depends_on_instances, SubscriptionInstanceRelationTable.depends_on)
+        .join(in_use_by_subscriptions, depends_on_instances.subscription)
+        .filter(depends_on_instances.subscription_id.in_(set(subscription_ids)))
+        .filter(in_use_by_instances.subscription_id != depends_on_instances.subscription_id)
+        .filter(in_use_by_subscriptions.status.in_(filter_statuses))
+    )
+
+    return _get_instance_relations(query_get_in_use_by_ids)
+
+
+async def _get_depends_on_instance_relations(
+    subscription_ids: list[UUID], filter_statuses: list[str]
+) -> list[Relation]:
+    """Get depends_on relations through subscription instance hierarchy."""
+    in_use_by_instances = aliased(SubscriptionInstanceTable)
+    depends_on_instances = aliased(SubscriptionInstanceTable)
+    depends_on_subscriptions = aliased(SubscriptionTable)
+
+    query_get_depends_on_ids = (
+        select(depends_on_instances.subscription_id, in_use_by_instances.subscription_id)
+        .distinct()
+        .join(depends_on_instances.subscription)
+        .join(depends_on_instances.in_use_by_block_relations)
+        .join(in_use_by_instances, SubscriptionInstanceRelationTable.in_use_by)
+        .join(depends_on_subscriptions, in_use_by_instances.subscription)
+        .filter(in_use_by_instances.subscription_id.in_(set(subscription_ids)))
+        .filter(depends_on_instances.subscription_id != in_use_by_instances.subscription_id)
+        .filter(depends_on_subscriptions.status.in_(filter_statuses))
+    )
+
+    return _get_instance_relations(query_get_depends_on_ids)
+
+
+def _get_resource_type_relations(resource_type_relations_query: Any) -> list[Relation]:
+    def to_relation(row: Row[Any]) -> Relation:
+        return Relation(UUID(row[0]), row[1])
+
+    return [to_relation(row) for row in db.session.execute(resource_type_relations_query)]
+
+
+async def _get_in_use_by_resource_type_relations(
+    subscription_ids: list[UUID], filter_statuses: list[str]
+) -> list[Relation]:
+    """Get in_use_by relations through resource types."""
+    logger.warning("Using legacy RELATION_RESOURCE_TYPES to find in_use_by subs")
+
+    in_use_by_subscriptions = aliased(SubscriptionTable)
+    depends_on_instance_values = aliased(SubscriptionInstanceValueTable)
+
+    # Convert UUIDs to string
+    unique_subscription_ids = set(map(str, subscription_ids))
+
+    query_get_in_use_by_ids = (
+        select(depends_on_instance_values.value, in_use_by_subscriptions.subscription_id)
+        .select_from(depends_on_instance_values)
+        .join(SubscriptionInstanceTable)
+        .join(in_use_by_subscriptions)
+        .join(ResourceTypeTable)
+        .filter(ResourceTypeTable.resource_type.in_(RELATION_RESOURCE_TYPES))
+        .filter(depends_on_instance_values.value.in_(unique_subscription_ids))
+        .filter(in_use_by_subscriptions.status.in_(filter_statuses))
+    )
+
+    return _get_resource_type_relations(query_get_in_use_by_ids)
+
+
+async def _get_depends_on_resource_type_relations(
+    subscription_ids: list[UUID], filter_statuses: list[str]
+) -> list[Relation]:
+    """Get depends_on relations through resource types."""
+    logger.warning("Using legacy RELATION_RESOURCE_TYPES to find depends_on subs")
+
+    depends_on_subscriptions = aliased(SubscriptionTable)
+    in_use_by_instances = aliased(SubscriptionInstanceTable)
+    in_use_by_instance_values = aliased(SubscriptionInstanceValueTable)
+
+    unique_subscription_ids = set(subscription_ids)
+
+    query_get_depends_on_ids = (
+        select(in_use_by_instance_values.value, in_use_by_instances.subscription_id)
+        .select_from(in_use_by_instance_values)
+        .join(in_use_by_instances)
+        .join(
+            depends_on_subscriptions,
+            in_use_by_instance_values.value == sa_cast(depends_on_subscriptions.subscription_id, SaText),
+        )
+        .join(ResourceTypeTable)
+        .filter(ResourceTypeTable.resource_type.in_(RELATION_RESOURCE_TYPES))
+        .filter(in_use_by_instances.subscription_id.in_(unique_subscription_ids))
+        .filter(depends_on_subscriptions.status.in_(filter_statuses))
+    )
+
+    return _get_resource_type_relations(query_get_depends_on_ids)
+
+
+async def _get_in_use_by_relations(subscription_ids: list[UUID], filter_statuses: list[str]) -> list[Relation]:
+    if RELATION_RESOURCE_TYPES:
+        # Find relations through resource types
+        resource_type_relations = await _get_in_use_by_resource_type_relations(subscription_ids, filter_statuses)
+    else:
+        resource_type_relations = []
+    # Find relations through instance hierarchy
+    instance_relations = await _get_in_use_by_instance_relations(subscription_ids, filter_statuses)
+    return list(chain(resource_type_relations, instance_relations))
+
+
+async def _get_depends_on_relations(subscription_ids: list[UUID], filter_statuses: list[str]) -> list[Relation]:
+    if RELATION_RESOURCE_TYPES:
+        # Find relations through resource types
+        resource_type_relations = await _get_depends_on_resource_type_relations(subscription_ids, filter_statuses)
+    else:
+        resource_type_relations = []
+    # Find relations through instance hierarchy
+    instance_relations = await _get_depends_on_instance_relations(subscription_ids, filter_statuses)
+    return list(chain(resource_type_relations, instance_relations))
+
+
+async def in_use_by_subs_loader(keys: list[tuple[UUID, list[str] | None]]) -> list[list[SubscriptionTable]]:
+    """GraphQL dataloader to efficiently get the in_use_by SubscriptionTables for multiple subscription_ids."""
+    subscription_ids = [key[0] for key in keys]
+    filter_statuses: list[str] = keys[0][1] or SubscriptionLifecycle.values()
+
+    in_use_by_relations = await _get_in_use_by_relations(subscription_ids, filter_statuses)
+
+    # Retrieve SubscriptionTable for all unique inuseby ids
+    unique_in_use_by_ids = {row.in_use_by_sub_id for row in in_use_by_relations}
+    _in_use_by_subs = db.session.execute(
+        select(SubscriptionTable).filter(SubscriptionTable.subscription_id.in_(unique_in_use_by_ids))
+    ).scalars()
+    in_use_by_subs = {subscription.subscription_id: subscription for subscription in _in_use_by_subs}
+
+    # group (more_itertools.bucket doesn't seem to work for tuple of uuids)
+    subscription_in_use_by_ids: dict[UUID, list[UUID]] = {}
+    for relation in in_use_by_relations:
+        subscription_in_use_by_ids.setdefault(relation.depends_on_sub_id, []).append(relation.in_use_by_sub_id)
+
+    def get_in_use_by_subs(depends_on_id: UUID) -> list[SubscriptionTable]:
+        in_use_by_ids = subscription_in_use_by_ids.get(depends_on_id, [])
+        return [in_use_by_sub for id_ in in_use_by_ids if (in_use_by_sub := in_use_by_subs.get(id_))]
+
+    # Important (as with any dataloader)
+    # Return the list of inuseby subs in the exact same order as the ids passed to this function
+    return [get_in_use_by_subs(subscription_id) for subscription_id in subscription_ids]
+
+
+async def depends_on_subs_loader(keys: list[tuple[UUID, list[str] | None]]) -> list[list[SubscriptionTable]]:
+    """GraphQL dataloader to efficiently get the depends_on SubscriptionTables for multiple subscription_ids."""
+    subscription_ids = [key[0] for key in keys]
+    filter_statuses: list[str] = keys[0][1] or SubscriptionLifecycle.values()
+
+    depends_on_relations = await _get_depends_on_relations(subscription_ids, filter_statuses)
+
+    # Retrieve SubscriptionTable for all unique dependson ids
+    unique_depends_on_ids = {row.depends_on_sub_id for row in depends_on_relations}
+    _depends_on_subs = db.session.execute(
+        select(SubscriptionTable).filter(SubscriptionTable.subscription_id.in_(unique_depends_on_ids))
+    ).scalars()
+    depends_on_subs = {subscription.subscription_id: subscription for subscription in _depends_on_subs}
+
+    # group (more_itertools.bucket doesn't seem to work for tuple of uuids)
+    subscription_depends_on_ids: dict[UUID, list[UUID]] = {}
+    for relation in depends_on_relations:
+        subscription_depends_on_ids.setdefault(relation.in_use_by_sub_id, []).append(relation.depends_on_sub_id)
+
+    def get_depends_on_subs(in_use_by_id: UUID) -> list[SubscriptionTable]:
+        depends_on_ids = subscription_depends_on_ids.get(in_use_by_id, [])
+        return [depends_on_sub for id_ in depends_on_ids if (depends_on_sub := depends_on_subs.get(id_))]
+
+    # Important (as with any dataloader)
+    # Return the list of dependson subs in the exact same order as the ids passed to this function
+    return [get_depends_on_subs(subscription_id) for subscription_id in subscription_ids]
+
+
+SubsLoaderType = DataLoader[tuple[UUID, list[str] | None], list[SubscriptionTable]]

orchestrator/graphql/loaders/subscriptions.py

@@ -10,14 +10,14 @@
 # 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 typing import cast
 from uuid import UUID
 
 import structlog
 from graphql import GraphQLError
 from pydantic.alias_generators import to_camel as to_lower_camel
 from sqlalchemy import Select, func, select
+from sqlalchemy.orm import contains_eager
 from strawberry.experimental.pydantic.conversion_types import StrawberryTypeFromPydantic
 
 from nwastdlib.asyncio import gather_nice
@@ -125,7 +125,16 @@ async def resolve_subscriptions(
         filter=pydantic_filter_by,
         query=query,
     )
-    stmt = select(SubscriptionTable).join(ProductTable)
+
+    stmt = (
+        select(SubscriptionTable)
+        .join(ProductTable)
+        .options(
+            # contains_eager() is needed because .join() does not eagerload, unlike options(joinedload())
+            # (and using joinedload() is not possible because of filter_subscriptions())
+            contains_eager(SubscriptionTable.product),
+        )
+    )
 
     stmt = filter_subscriptions(stmt, pydantic_filter_by, _error_handler)
     if query is not None:

orchestrator/graphql/resolvers/subscription.py

@@ -13,7 +13,7 @@
 from collections.abc import Callable, Iterable
 from http import HTTPStatus
 from pathlib import Path
-from typing import Any, Coroutine
+from typing import Any, Coroutine, Protocol
 
 import strawberry
 import structlog
@@ -134,17 +134,26 @@ def process_errors(
                 StrawberryLogger.error(error, execution_context)
 
 
-def get_context(
+class ContextGetterFactory(Protocol):
+    def __call__(
+        self,
+        auth_manager: AuthManager,
+        graphql_models: StrawberryModelType,
+        broadcast_thread: ProcessDataBroadcastThread | None = None,
+    ) -> Callable[[], Coroutine[Any, Any, OrchestratorContext]]: ...
+
+
+def default_context_getter(
     auth_manager: AuthManager,
     graphql_models: StrawberryModelType,
     broadcast_thread: ProcessDataBroadcastThread | None = None,
 ) -> Callable[[], Coroutine[Any, Any, OrchestratorContext]]:
-    async def _get_context() -> OrchestratorContext:
+    async def context_getter() -> OrchestratorContext:
         return OrchestratorContext(
             auth_manager=auth_manager, graphql_models=graphql_models, broadcast_thread=broadcast_thread
         )
 
-    return _get_context
+    return context_getter
 
 
 def get_extensions(mutation: Any, query: Any) -> Iterable[type[SchemaExtension]]:
@@ -169,6 +178,7 @@ def create_graphql_router(
     graphql_models: StrawberryModelType | None = None,
     scalar_overrides: ScalarOverrideType | None = None,
     extensions: list | None = None,
+    custom_context_getter: ContextGetterFactory | None = None,
 ) -> OrchestratorGraphqlRouter:
     scalar_overrides = scalar_overrides if scalar_overrides else dict(SCALAR_OVERRIDES)
     models = graphql_models if graphql_models else dict(DEFAULT_GRAPHQL_MODELS)
@@ -190,8 +200,9 @@ def create_graphql_router(
         scalar_overrides=scalar_overrides,
     )
 
+    context_getter_factory = custom_context_getter or default_context_getter
     return OrchestratorGraphqlRouter(
         schema,
-        context_getter=get_context(auth_manager, models, broadcast_thread),
+        context_getter=context_getter_factory(auth_manager, models, broadcast_thread),
         graphiql=app_settings.SERVE_GRAPHQL_UI,
     )