Merge branch 'deploy/hammer' into hammer/generic-alerts

Signed-off-by: Aaron Chong <aaronchongth@gmail.com>

packages/api-client/lib/version.ts

@@ -3,6 +3,6 @@ import { version as rmfModelVer } from 'rmf-models';
 
 export const version = {
   rmfModels: rmfModelVer,
-  rmfServer: 'eb83039e43d7caee62ea130adf1f73b24d8465e6',
+  rmfServer: 'cb1b563ab8f827756f9a453ec944362d92a3ab57',
   openapiGenerator: '6.2.1',
 };

packages/api-server/.pylintrc

@@ -152,6 +152,7 @@ disable=print-statement,
         logging-fstring-interpolation,
         line-too-long,
         too-many-lines,
+        no-self-use,
 
 # fastapi heavily uses singletons, disallowing global statements just ends up with things
 # using nonlocal as singletons which are functionally the same as globals.

packages/api-server/README.md

@@ -252,18 +252,24 @@ Restart the `api-server` and the changes to the databse should be reflected.
 ### Running unit tests
 
 ```bash
-npm test
+pnpm test
+```
+
+By default in-memory sqlite database is used for testing, to test on another database, set the `RMF_API_SERVER_TEST_DB_URL` environment variable.
+
+```bash
+RMF_API_SERVER_TEST_DB_URL=<db_url> pnpm test
 ```
 
 ### Collecting code coverage
 
 ```bash
-npm run test:cov
+pnpm run test:cov
 ```
 
 Generate coverage report
 ```bash
-npm run test:report
+pnpm run test:report
 ```
 
 ## Live reload

packages/api-server/api_server/dependencies.py

@@ -20,7 +20,11 @@ def pagination_query(
 ) -> Pagination:
     limit = limit or 100
     offset = offset or 0
-    return Pagination(limit=limit, offset=offset, order_by=order_by)
+    return Pagination(
+        limit=limit,
+        offset=offset,
+        order_by=order_by.split(",") if order_by else [],
+    )
 
 
 # hacky way to get the sio user

packages/api-server/api_server/models/__init__.py

@@ -7,6 +7,7 @@
 from .doors import *
 from .health import *
 from .ingestors import *
+from .labels import *
 from .lifts import *
 from .pagination import *
 from .rmf_api.activity_discovery_request import ActivityDiscoveryRequest

packages/api-server/api_server/models/labels.py

@@ -0,0 +1,25 @@
+from typing import Sequence
+
+from pydantic import BaseModel
+
+
+class Labels(BaseModel):
+    """
+    Labels for a resource.
+    """
+
+    __root__: dict[str, str]
+
+    @staticmethod
+    def _parse_label(s: str) -> tuple[str, str]:
+        sep = s.find("=")
+        if sep == -1:
+            return s, ""
+        return s[:sep], s[sep + 1 :]
+
+    @staticmethod
+    def from_strings(labels: Sequence[str]) -> "Labels":
+        return Labels(__root__=dict(Labels._parse_label(s) for s in labels))
+
+    def to_strings(self) -> list[str]:
+        return [f"{k}={v}" for k, v in self.__root__.items()]

packages/api-server/api_server/models/pagination.py

@@ -1,9 +1,7 @@
-from typing import Optional
-
 from pydantic import BaseModel
 
 
 class Pagination(BaseModel):
     limit: int
     offset: int
-    order_by: Optional[str]
+    order_by: list[str]

packages/api-server/api_server/models/rmf_api/task_request.py

@@ -35,7 +35,3 @@ class TaskRequest(BaseModel):
         None,
         description="(Optional) The name of the fleet that should perform this task. If specified, other fleets will not bid for this task.",
     )
-    unix_millis_warn_time: Optional[int] = Field(
-        None,
-        description="(Optional) The time at which a warning will be issued if the estimated completion time is later than expected",
-    )

packages/api-server/api_server/models/rmf_api/task_state.py

@@ -212,4 +212,3 @@ class TaskState(BaseModel):
         None,
         description="If the task was killed, this will describe information about the request.",
     )
-    unix_millis_warn_time: Optional[int] = None

packages/api-server/api_server/models/task_booking_label.py

@@ -3,30 +3,6 @@
 import pydantic
 from pydantic import BaseModel
 
-# NOTE: This label model needs to exactly match the fields that are defined and
-# populated by the dashboard. Any changes to either side will require syncing.
-
-
-class TaskBookingLabelDescription(BaseModel):
-    """
-    This description holds several fields that could be useful for frontend
-    dashboards when dispatching a task, to then be identified or rendered
-    accordingly back on the same frontend.
-    """
-
-    task_definition_id: str
-    unix_millis_warn_time: Optional[int]
-    pickup: Optional[str]
-    destination: Optional[str]
-    cart_id: Optional[str]
-
-    @staticmethod
-    def from_json_string(json_str: str) -> Optional["TaskBookingLabelDescription"]:
-        try:
-            return TaskBookingLabelDescription.parse_raw(json_str)
-        except pydantic.error_wrappers.ValidationError:
-            return None
-
 
 class TaskBookingLabel(BaseModel):
     """
@@ -36,7 +12,7 @@ class TaskBookingLabel(BaseModel):
     needed for any frontends.
     """
 
-    description: TaskBookingLabelDescription
+    description: dict[str, str | int | float]
 
     @staticmethod
     def from_json_string(json_str: str) -> Optional["TaskBookingLabel"]:

packages/api-server/api_server/models/tortoise_models/tasks.py

@@ -28,10 +28,14 @@ class TaskState(Model):
     status = CharField(255, null=True, index=True)
     unix_millis_request_time = DatetimeField(null=True, index=True)
     requester = CharField(255, null=True, index=True)
-    unix_millis_warn_time = DatetimeField(null=True, index=True)
-    pickup = CharField(255, null=True, index=True)
-    destination = CharField(255, null=True, index=True)
+    _do_not_use_pickup = CharField(255, null=True, index=True, source_field="pickup")
+    _do_not_use_destination = CharField(
+        255, null=True, index=True, source_field="destination"
+    )
     labels = ReverseRelation["TaskLabel"]
+    _do_not_use_unix_millis_warn_time = DatetimeField(
+        null=True, index=True, source_field="unix_millis_warn_time"
+    )
 
 
 class TaskLabel(Model):

packages/api-server/api_server/query.py

@@ -1,5 +1,3 @@
-from typing import Dict, Optional
-
 from tortoise.queryset import MODEL, QuerySet
 
 from api_server.models.pagination import Pagination
@@ -8,25 +6,10 @@
 def add_pagination(
     query: QuerySet[MODEL],
     pagination: Pagination,
-    field_mappings: Optional[Dict[str, str]] = None,
 ) -> QuerySet[MODEL]:
-    """
-    Adds pagination and ordering to a query.
-
-    :param field_mapping: A dict mapping the order fields to the fields used to build the
-        query. e.g. a url of `?order_by=order_field` and a field mapping of `{"order_field": "db_field"}`
-        will order the query result according to `db_field`.
-    """
-    field_mappings = field_mappings or {}
-    query = query.limit(pagination.limit).offset(pagination.offset)
-    if pagination.order_by is not None:
-        order_fields = []
-        order_values = pagination.order_by.split(",")
-        for v in order_values:
-            if v[0] in ["-", "+"]:
-                stripped = v[1:]
-                order_fields.append(v[0] + field_mappings.get(stripped, stripped))
-            else:
-                order_fields.append(field_mappings.get(v, v))
-        query = query.order_by(*order_fields)
-    return query
+    """Adds pagination and ordering to a query"""
+    return (
+        query.limit(pagination.limit)
+        .offset(pagination.offset)
+        .order_by(*pagination.order_by)
+    )

packages/api-server/api_server/repositories/tasks.py

@@ -1,18 +1,18 @@
 from datetime import datetime
 from typing import Dict, List, Optional, Sequence, Tuple, cast
 
+import tortoise.functions as tfuncs
 from fastapi import Depends, HTTPException
 from tortoise.exceptions import FieldError, IntegrityError
+from tortoise.expressions import Expression, Q
 from tortoise.query_utils import Prefetch
-from tortoise.queryset import QuerySet
 from tortoise.transactions import in_transaction
 
 from api_server.authenticator import user_dep
 from api_server.logging import LoggerAdapter, get_logger
+from api_server.models import Labels, LogEntry, Pagination, Phases
+from api_server.models import Status as TaskStatus
 from api_server.models import (
-    LogEntry,
-    Pagination,
-    Phases,
     TaskBookingLabel,
     TaskEventLog,
     TaskRequest,
@@ -22,7 +22,6 @@
 from api_server.models import tortoise_models as ttm
 from api_server.models.tortoise_models import TaskRequest as DbTaskRequest
 from api_server.models.tortoise_models import TaskState as DbTaskState
-from api_server.query import add_pagination
 
 
 class TaskRepository:
@@ -113,31 +112,115 @@ async def save_task_state(self, task_state: TaskState) -> None:
 
         # Here we generate the labels required for server-side sorting and
         # filtering.
-        if booking_label.description.pickup is not None:
-            await ttm.TaskLabel.create(
-                state=state,
-                label_name="pickup",
-                label_value_str=booking_label.description.pickup,
-            )
-        if booking_label.description.destination is not None:
-            await ttm.TaskLabel.create(
-                state=state,
-                label_name="destination",
-                label_value_str=booking_label.description.destination,
-            )
-        if booking_label.description.unix_millis_warn_time is not None:
-            await ttm.TaskLabel.create(
-                state=state,
-                label_name="unix_millis_warn_time",
-                label_value_num=booking_label.description.unix_millis_warn_time,
-            )
+        async with in_transaction():
+            for k, v in booking_label.description.items():
+                if isinstance(v, str):
+                    await ttm.TaskLabel.create(
+                        state=state, label_name=k, label_value_str=v
+                    )
+                elif isinstance(v, int):
+                    await ttm.TaskLabel.create(
+                        state=state,
+                        label_name=k,
+                        label_value_num=v,
+                        label_value_float=v,  # also store float to make querying easier
+                    )
+                elif isinstance(v, float):
+                    exact_val = int(v) if v.is_integer else None
+                    await ttm.TaskLabel.create(
+                        state=state,
+                        label_name=k,
+                        label_value_float=v,
+                        label_value_num=exact_val,  # also store int to make querying easier
+                    )
 
     async def query_task_states(
-        self, query: QuerySet[DbTaskState], pagination: Optional[Pagination] = None
+        self,
+        task_id: list[str] | None = None,
+        category: list[str] | None = None,
+        assigned_to: list[str] | None = None,
+        start_time_between: tuple[datetime, datetime] | None = None,
+        finish_time_between: tuple[datetime, datetime] | None = None,
+        request_time_between: tuple[datetime, datetime] | None = None,
+        requester: list[str] | None = None,
+        status: list[str] | None = None,
+        label: Labels | None = None,
+        pagination: Optional[Pagination] = None,
     ) -> List[TaskState]:
+        filters = {}
+        if task_id is not None:
+            filters["id___in"] = task_id
+        if category is not None:
+            filters["category__in"] = category
+        if assigned_to is not None:
+            filters["assigned_to__in"] = assigned_to
+        if start_time_between is not None:
+            filters["unix_millis_start_time__gte"] = start_time_between[0]
+            filters["unix_millis_start_time__lte"] = start_time_between[1]
+        if finish_time_between is not None:
+            filters["unix_millis_finish_time__gte"] = finish_time_between[0]
+            filters["unix_millis_finish_time__lte"] = finish_time_between[1]
+        if request_time_between is not None:
+            filters["unix_millis_request_time__gte"] = request_time_between[0]
+            filters["unix_millis_request_time__lte"] = request_time_between[1]
+        if requester is not None:
+            filters["requester__in"] = requester
+        if status is not None:
+            valid_values = [member.value for member in TaskStatus]
+            filters["status__in"] = []
+            for status_string in status:
+                if status_string not in valid_values:
+                    continue
+                filters["status__in"].append(TaskStatus(status_string))
+        query = DbTaskState.filter(**filters)
+
+        need_group_by = False
+        label_filters = {}
+        if label is not None:
+            label_filters.update(
+                {
+                    f"label_filter_{k}": tfuncs.Count(
+                        "id_",
+                        _filter=Q(labels__label_name=k, labels__label_value_str=v),
+                    )
+                    for k, v in label.__root__.items()
+                }
+            )
+
+        if len(label_filters) > 0:
+            filter_gt = {f"{f}__gt": 0 for f in label_filters}
+            query = query.annotate(**label_filters).filter(**filter_gt)
+            need_group_by = True
+
+        if pagination:
+            order_fields: list[str] = []
+            annotations: dict[str, Expression] = {}
+            # add annotations required for sorting by labels
+            for f in pagination.order_by:
+                order_prefix = f[0] if f[0] == "-" else ""
+                order_field = f[1:] if order_prefix == "-" else f
+                if order_field.startswith("label="):
+                    f = order_field[6:]
+                    annotations[f"label_sort_{f}"] = tfuncs.Max(
+                        "labels__label_value_str",
+                        _filter=Q(labels__label_name=f),
+                    )
+                    order_field = f"label_sort_{f}"
+
+                order_fields.append(order_prefix + order_field)
+
+            query = (
+                query.annotate(**annotations)
+                .limit(pagination.limit)
+                .offset(pagination.offset)
+                .order_by(*order_fields)
+            )
+            need_group_by = True
+
+        if need_group_by:
+            query = query.group_by("id_", "labels__state_id")
+
         try:
-            if pagination:
-                query = add_pagination(query, pagination)
             # TODO: enforce with authz
             results = await query.values_list("data", flat=True)
             return [TaskState(**r) for r in results]

packages/api-server/api_server/routes/tasks/scheduled_tasks.py

@@ -182,7 +182,7 @@ async def get_scheduled_tasks(
         .offset(pagination.offset)
     )
     if pagination.order_by:
-        q.order_by(*pagination.order_by.split(","))
+        q.order_by(*pagination.order_by)
     return await ttm.ScheduledTaskPydanticList.from_queryset(q)