Merge branch 'main' into tiny-cleanup

Author: induane

cli/entry_points/init.py

@@ -168,7 +168,7 @@ def initialize_database(self) -> None:
                         description="This key is utilized by the product demo",
                     )
 
-                    action_args = {"from_address": "[email protected]", "recipients": [], "template": "NotifyTemplate"}
+                    action_args = {"recipients": [], "template": "NotifyTemplate"}
                     action = Action.create(
                         name="Send Email", action_impl="SEND_EMAIL", company=company, action_args=action_args
                     )

common/actions/send_email_action.py

@@ -19,7 +19,7 @@
 
 
 class SendEmailAction(BaseAction):
-    required_arguments = {"from_address", "recipients", "template"}
+    required_arguments = {"recipients", "template"}
     requires_action_template = True
 
     def _run(self, event: EVENT_TYPE, rule: Rule, journey_id: Optional[UUID]) -> ActionResult:
@@ -29,8 +29,8 @@ def _run(self, event: EVENT_TYPE, rule: Rule, journey_id: Optional[UUID]) -> Act
             return ActionResult(False, None, e)
         try:
             response = EmailService.send_email(
-                self.arguments["smtp_config"],
-                self.arguments["from_address"],
+                self.arguments.get("smtp_config", {}),
+                self.arguments.get("from_address"),
                 self.arguments["recipients"],
                 self.arguments["template"],
                 context,

common/email/email_service.py

@@ -17,12 +17,13 @@ class EmailService:
     @staticmethod
     def send_email(
         smtp_config: dict,
-        from_address: str,
+        from_address: str | None,
         recipients: list[str],
         template_name: str,
         template_context_vars: Mapping,
     ) -> dict:
         try:
+            from_address = from_address or settings.SMTP["from_address"]
             content, subject = HandlebarsEmailRenderer.render(template_name, template_context_vars)
             context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
             message = MIMEMultipart("alternative")

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

@@ -9,9 +9,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")
@@ -51,9 +53,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

@@ -51,16 +51,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,34 +55,36 @@ 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((TestOutcome.name ** f"%{rules.search}%") | (TestOutcome.description ** f"%{rules.search}%"))
+            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/schemas/action_schemas.py

@@ -16,7 +16,7 @@ class SMTPConfigSchema(Schema):
 
 
 class EmailActionArgsSchema(Schema):
-    from_address = Str(validate=not_empty(max=255))
+    from_address = Str(required=False, validate=not_empty(max=255))
     template = Str(validate=not_empty(max=255))
     recipients = List(Email(), required=True, validate=Length(max=50))
     smtp_config = Nested(SMTPConfigSchema(), required=False)