Merge branch 'events-list-performance' into 'enterprise'

refactor(be): remove coalesce when ordering events list

See merge request dkinternal/observability/dataops-observability!24

Author: ci bot

common/entities/event.py

@@ -2,7 +2,7 @@
 
 from enum import Enum, IntEnum
 
-from peewee import ForeignKeyField, Node, fn
+from peewee import ForeignKeyField
 from playhouse.mysql_ext import JSONField
 
 from common.peewee_extensions.fields import EnumIntField, EnumStrField, UTCTimestampField
@@ -44,10 +44,6 @@ class EventEntity(BaseEntity):
     instance_set = ForeignKeyField(InstanceSet, null=True, backref="events", on_delete="SET NULL")
     v2_payload = JSONField(null=False)
 
-    @classmethod
-    def timestamp_coalesce(cls) -> Node:
-        return fn.COALESCE(cls.timestamp, cls.created_timestamp)
-
     @property
     def components(self) -> list[Component]:
         return [self.component] if self.component else []

common/entity_services/company_service.py

@@ -26,7 +26,7 @@ def get_organizations_with_rules(company_id: str, rules: ListRules) -> Page[Orga
 
     @staticmethod
     def get_users_with_rules(company_id: str, rules: ListRules) -> Page[User]:
-        query = User.select().join(Company).where(Company.id == company_id)
+        query = User.select().join(Company).where(User.primary_company == company_id)
         return Page[User].get_paginated_results(query, User.name, rules)
 
     @staticmethod

common/entity_services/event_service.py

@@ -37,11 +37,12 @@ def get_events_with_rules(*, rules: ListRules, filters: ProjectEventFilters) ->
         filter_list: list[object] = [EventEntity.project << filters.project_ids]
 
         if filters.instance_ids or filters.journey_ids:
-            query = query.join(InstanceSet).join(InstancesInstanceSets).join(Instance).switch(EventEntity)
+            instance_set_subquery = InstanceSet.select(InstanceSet.id).join(InstancesInstanceSets).join(Instance)
             if filters.journey_ids:
-                filter_list.append(Instance.journey << filters.journey_ids)
+                instance_set_subquery = instance_set_subquery.where(Instance.journey << filters.journey_ids)
             if filters.instance_ids:
-                filter_list.append(Instance.id << filters.instance_ids)
+                instance_set_subquery = instance_set_subquery.where(Instance.id << filters.instance_ids)
+            query = query.where(EventEntity.instance_set.in_(instance_set_subquery))
         if filters.event_types:
             filter_list.append(EventEntity.type << filters.event_types)
         if filters.component_ids:
@@ -53,16 +54,12 @@ def get_events_with_rules(*, rules: ListRules, filters: ProjectEventFilters) ->
         if filters.task_ids:
             filter_list.append(EventEntity.task << filters.task_ids)
         if filters.date_range_start:
-            filter_list.append(
-                EventEntity.timestamp_coalesce() >= EventEntity.timestamp.db_value(filters.date_range_start)
-            )
+            filter_list.append(EventEntity.timestamp >= filters.date_range_start)
         if filters.date_range_end:
-            filter_list.append(
-                EventEntity.timestamp_coalesce() < EventEntity.timestamp.db_value(filters.date_range_end)
-            )
+            filter_list.append(EventEntity.timestamp < filters.date_range_end)
 
         query = query.where(*filter_list)
-        page = Page[EventEntity].get_paginated_results(query, EventEntity.timestamp_coalesce(), rules)
+        page = Page[EventEntity].get_paginated_results(query, EventEntity.timestamp, rules)
 
         # Using a single query to fetch the Instance and Journey data
         instance_set_ids = {e.instance_set_id for e in page.results}

common/entity_services/helpers/list_rules.py

@@ -8,9 +8,11 @@
 
 from marshmallow import EXCLUDE, Schema
 from marshmallow.fields import Enum, Int
-from peewee import Field, Ordering, Select
+from peewee import JOIN, Field, Ordering, Select
 from werkzeug.datastructures import MultiDict
 
+from common.entities import BaseEntity
+
 DEFAULT_PAGE = 1
 DEFAULT_COUNT = 10
 T = TypeVar("T")
@@ -50,9 +52,24 @@ def __len__(self) -> int:
 
     @classmethod
     def get_paginated_results(cls, query: Select, order_by: Field, list_rules: ListRules) -> Page[T]:
+        model: BaseEntity = query.model
         ordering = list_rules.order_by_field(order_by)
-        paginated_query = query.order_by(ordering).paginate(list_rules.page, list_rules.count)
-        return cls(results=list(paginated_query), total=query.count())
+
+        paginated_subquery: Select = (
+            model.select(model.id)
+            .where(query._where)
+            .order_by(ordering)
+            .paginate(list_rules.page, list_rules.count)
+            .alias("results")
+        )
+
+        results_query = query.clone()
+        results_query._where = None
+        results_query = results_query.join(
+            paginated_subquery, join_type=JOIN.INNER, on=(model.id == paginated_subquery.c.id)
+        ).order_by(ordering)
+
+        return cls(results=list(results_query), total=paginated_subquery.count(clear_limit=True))
 
 
 @dataclass

common/entity_services/journey_service.py

@@ -52,16 +52,18 @@ def get_action_by_implementation(journey_id: UUID, action_impl: str) -> Optional
 
     @staticmethod
     def get_components_with_rules(journey_id: str, rules: ListRules, filters: ComponentFilters) -> Page[Component]:
-        query = JourneyDagEdge.journey_id == journey_id
+        join_on = (JourneyDagEdge.journey_id == journey_id) & (
+            (JourneyDagEdge.left == Component.id) | (JourneyDagEdge.right == Component.id)
+        )
+        query = Component.select(Component).join(JourneyDagEdge, on=join_on)
         if rules.search is not None:
-            query &= Component.key ** f"%{rules.search}%"
+            query = query.where(Component.key ** f"%{rules.search}%")
         if filters.component_types:
-            query &= Component.type.in_(filters.component_types)
+            query = query.where(Component.type.in_(filters.component_types))
         if filters.tools:
-            query &= Component.tool.in_(filters.tools)
-        join_on = (JourneyDagEdge.left == Component.id) | (JourneyDagEdge.right == Component.id)
-        query = Component.select(Component).join(JourneyDagEdge, on=join_on).where(query).distinct()
-        return Page[Component].get_paginated_results(query, Component.key, rules)
+            query = query.where(Component.tool.in_(filters.tools))
+
+        return Page[Component].get_paginated_results(query.distinct(), Component.key, rules)
 
     @staticmethod
     def get_upstream_nodes(journey: Journey, component_id: UUID) -> set:

common/entity_services/project_service.py

@@ -55,36 +55,37 @@ class ProjectService:
     def get_test_outcomes_with_rules(
         project: Project, rules: ListRules, filters: TestOutcomeItemFilters
     ) -> Page[TestOutcome]:
-        clauses = [TestOutcome.component.in_(project.components)]
+        query = TestOutcome.select(TestOutcome).where(TestOutcome.component.in_(project.components))
         if rules.search is not None:
-            clauses.append(
+            query = query.where(
                 ((TestOutcome.name ** f"%{rules.search}%") | (TestOutcome.description ** f"%{rules.search}%"))
             )
+
         if filters:
             if statuses := filters.statuses:
-                clauses.append(TestOutcome.status.in_(statuses))
+                query = query.where(TestOutcome.status.in_(statuses))
             if component_ids := filters.component_ids:
-                clauses.append(TestOutcome.component << component_ids)
+                query = query.where(TestOutcome.component << component_ids)
             if start_range_begin := filters.start_range_begin:
-                clauses.append(TestOutcome.start_time >= start_range_begin)
+                query = query.where(TestOutcome.start_time >= start_range_begin)
             if start_range_end := filters.start_range_end:
-                clauses.append(TestOutcome.start_time < start_range_end)
+                query = query.where(TestOutcome.start_time < start_range_end)
             if end_range_begin := filters.end_range_begin:
-                clauses.append(TestOutcome.end_time >= end_range_begin)
+                query = query.where(TestOutcome.end_time >= end_range_begin)
             if end_range_end := filters.end_range_end:
-                clauses.append(TestOutcome.end_time < end_range_end)
+                query = query.where(TestOutcome.end_time < end_range_end)
             if run_ids := filters.run_ids:
-                clauses.append(TestOutcome.run.in_(run_ids))
+                query = query.where(TestOutcome.run.in_(run_ids))
             if instance_ids := filters.instance_ids:
-                clauses.append(InstancesInstanceSets.instance.in_(instance_ids))
+                instance_set_subquery = (
+                    InstanceSet.select(InstanceSet.id)
+                    .join(InstancesInstanceSets)
+                    .where(InstancesInstanceSets.instance.in_(instance_ids))
+                )
+                query = query.where(TestOutcome.instance_set.in_(instance_set_subquery))
             if key := filters.key:
-                clauses.append(TestOutcome.key == key)
+                query = query.where(TestOutcome.key == key)
 
-        # If filtering on instance_ids we need to join the InstanceSet tables
-        if filters and filters.instance_ids:
-            query = TestOutcome.select(TestOutcome).join(InstanceSet).join(InstancesInstanceSets).where(*clauses)
-        else:
-            query = TestOutcome.select(TestOutcome).where(*clauses)
         return Page[TestOutcome].get_paginated_results(query, TestOutcome.start_time, rules)
 
     @staticmethod

common/tests/unit/entity_services/helpers/test_list_rules.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import ANY, Mock, MagicMock
 
 import pytest
 from werkzeug.datastructures import MultiDict
@@ -20,21 +20,29 @@ def mocked_entity_class():
 
 @pytest.fixture()
 def mocked_query():
-    # .paginate() (and other query methods in general) doesn't return a list until iterated on, but this is good enough for the test
-    paginate = Mock(return_value=list())
     query = Mock()
-    query.paginate = paginate
-    query.order_by = Mock(return_value=query)
+    subquery = Mock()
+    results_query = MagicMock()
+
+    query.model.select.return_value = subquery
+    query.clone = Mock(return_value=results_query)
+
+    subquery.where.return_value = subquery
+    subquery.order_by.return_value = subquery
+    subquery.paginate.return_value = subquery
+    subquery.alias.return_value = subquery
+
+    results_query.join.return_value = results_query
+    results_query.order_by.return_value = results_query
+
     yield query
 
 
 @pytest.fixture()
-def mocked_query_with_data():
-    # .paginate() (and other query methods in general) doesn't return a list until iterated on, but this is good enough for the test
-    paginate = Mock(return_value=[Mock(), Mock()])
-    query = Mock()
-    query.paginate = paginate
-    query.order_by = Mock(return_value=query)
+def mocked_query_with_data(mocked_query):
+    query = mocked_query
+    results_query = query.clone.return_value
+    results_query.__iter__.return_value = [Mock(), Mock()]
     yield query
 
 
@@ -80,11 +88,24 @@ def test_from_params_extras():
 
 @pytest.mark.unit
 def test_page_get_paginated_results(mocked_field, mocked_entity_class, mocked_query):
+    mocked_subquery = mocked_query.model.select.return_value
+    mocked_results_query = mocked_query.clone.return_value
     rules = ListRules.from_params(MultiDict())
+
     result = Page[mocked_entity_class].get_paginated_results(mocked_query, mocked_field, rules)
+
     mocked_field.asc.assert_called_once()
-    mocked_query.paginate.assert_called_once()
-    mocked_query.count.assert_called_once()
+
+    mocked_subquery.where.assert_called_once_with(mocked_query._where)
+    mocked_subquery.order_by.assert_called_once_with(rules.order_by_field(mocked_field))
+    mocked_subquery.paginate.assert_called_once_with(rules.page, rules.count)
+    mocked_subquery.alias.assert_called_once()
+    mocked_subquery.count.assert_called_once_with(clear_limit=True)
+
+    assert mocked_results_query._where is None
+    mocked_results_query.join.assert_called_once_with(mocked_subquery, join_type=ANY, on=ANY)
+    mocked_results_query.order_by.assert_called_once_with(rules.order_by_field(mocked_field))
+
     assert type(result) == Page
 
 
@@ -102,6 +123,6 @@ def test_page_len(mocked_field, mocked_entity_class, mocked_query_with_data):
 def test_page_iter(mocked_field, mocked_entity_class, mocked_query_with_data):
     rules = ListRules.from_params(MultiDict())
     result = Page[mocked_entity_class].get_paginated_results(mocked_query_with_data, mocked_field, rules)
-    count = len(mocked_query_with_data.paginate())
+    count = len(list(mocked_query_with_data.clone.return_value))
     r = [r for r in result]
     assert len(r) == count

deploy/charts/observability-app/templates/observability-ui.yaml

@@ -51,6 +51,11 @@ spec:
               protocol: TCP
           resources:
             {{- toYaml .Values.observability_ui.resources | nindent 12 }}
+          {{- if .Values.observability_api.hostname }}
+          env:
+            - name: OBSERVABILITY_API_HOSTNAME
+              value: {{ tpl .Values.observability_api.hostname . | quote }}
+          {{- end }}
           {{- if .Values.observability_ui.environmentJson }}
           volumeMounts:
             - mountPath: /observability_ui/shell/environments

testlib/fixtures/entities.py

@@ -621,6 +621,7 @@ def event_entity_2(test_db, dataset):
         version=EventVersion.V2,
         type=ApiEventType.DATASET_OPERATION,
         created_timestamp=datetime(2024, 1, 20, 9, 55, 0, tzinfo=timezone.utc),
+        timestamp=datetime(2024, 1, 20, 9, 55, 0, tzinfo=timezone.utc),
         project=dataset.project_id,
         component=dataset,
         v2_payload={},