Support workers filtering with tags

apiserver/apimodels/workers.py

@@ -96,6 +96,7 @@ class WorkerResponseEntry(WorkerEntry):
 
 class GetAllRequest(Base):
     last_seen = IntField(default=3600)
+    tags = ListField(str)
 
 
 class GetAllResponse(Base):

apiserver/bll/workers/__init__.py

@@ -76,7 +76,7 @@ class WorkerBLL:
                 raise bad_request.InvalidUserId(**query)
             company = Company.objects(id=company_id).only("id", "name").first()
             if not company:
-                raise server_error.InternalError("invalid company", company=company_id)
+                raise bad_request.InvalidId("invalid company", company=company_id)
 
             queue_objs = Queue.objects(company=company_id, id__in=queues).only("id")
             if len(queue_objs) < len(queues):
@@ -189,7 +189,10 @@ class WorkerBLL:
             self._save_worker(entry)
 
     def get_all(
-        self, company_id: str, last_seen: Optional[int] = None
+        self,
+        company_id: str,
+        last_seen: Optional[int] = None,
+        tags: Sequence[str] = None,
     ) -> Sequence[WorkerEntry]:
         """
         Get all the company workers that were active during the last_seen period
@@ -210,16 +213,26 @@ class WorkerBLL:
                 if w.last_activity_time.replace(tzinfo=None) >= ref_time
             ]
 
+        if tags:
+            include = {t for t in tags if not t.startswith("-")}
+            exclude = {t[1:] for t in tags if t.startswith("-")}
+            workers = [
+                w
+                for w in workers
+                if (not include or any(t in include for t in w.tags))
+                and (not exclude or all(t not in exclude for t in w.tags))
+            ]
+
         return workers
 
     def get_all_with_projection(
-        self, company_id: str, last_seen: int
+        self, company_id: str, last_seen: int, tags: Sequence[str] = None
     ) -> Sequence[WorkerResponseEntry]:
 
         helpers = list(
             map(
                 WorkerConversionHelper.from_worker_entry,
-                self.get_all(company_id=company_id, last_seen=last_seen),
+                self.get_all(company_id=company_id, last_seen=last_seen, tags=tags),
             )
         )
 

apiserver/schema/services/workers.conf

@@ -1,6 +1,5 @@
-{
-    _description: "Provides an API for worker machines, allowing workers to report status and get tasks for execution"
-    _definitions {
+_description: "Provides an API for worker machines, allowing workers to report status and get tasks for execution"
+_definitions {
     metrics_category {
         type: object
         properties {
@@ -263,8 +262,8 @@
             }
         }
     }
-    }
-    get_all {
+}
+get_all {
     "2.4" {
         description: "Returns information on all registered workers."
         request {
@@ -288,8 +287,15 @@
             }
         }
     }
+    "999.0": ${get_all."2.4"} {
+        request.properties.tags {
+            description: The list of allowed worker tags. Prepend tag value with '-' in order to exclude
+            type: array
+            items { type: string }
         }
-    register {
+    }
+}
+register {
     "2.4" {
         description: "Register a worker in the system. Called by the Worker Daemon."
         request {
@@ -322,8 +328,8 @@
             properties {}
         }
     }
-    }
-    unregister {
+}
+unregister {
     "2.4" {
         description: "Unregister a worker in the system. Called by the Worker Daemon."
         request {
@@ -341,8 +347,8 @@
             properties {}
         }
     }
-    }
-    status_report {
+}
+status_report {
     "2.4" {
         description: "Called periodically by the worker daemon to report machine status"
         request {
@@ -389,8 +395,8 @@
             properties {}
         }
     }
-    }
-    get_metric_keys {
+}
+get_metric_keys {
     "2.4" {
         description: "Returns worker statistics metric keys grouped by categories."
         request {
@@ -414,8 +420,8 @@
             }
         }
     }
-    }
-    get_stats {
+}
+get_stats {
     "2.4" {
         description: "Returns statistics for the selected workers and time range aggregated by date intervals."
         request {
@@ -462,8 +468,8 @@
             }
         }
     }
-    }
-    get_activity_report {
+}
+get_activity_report {
     "2.4" {
         description: "Returns count of active company workers in the selected time range."
         request {
@@ -498,5 +504,4 @@
             }
         }
     }
-    }
 }

apiserver/services/workers.py

@@ -41,7 +41,9 @@ worker_bll = WorkerBLL()
 )
 def get_all(call: APICall, company_id: str, request: GetAllRequest):
     call.result.data_model = GetAllResponse(
-        workers=worker_bll.get_all_with_projection(company_id, request.last_seen)
+        workers=worker_bll.get_all_with_projection(
+            company_id, request.last_seen, tags=request.tags
+        )
     )
 
 
@@ -72,7 +74,9 @@ def unregister(call: APICall, company_id, req_model: WorkerRequest):
     worker_bll.unregister_worker(company_id, call.identity.user, req_model.worker)
 
 
-@endpoint("workers.status_report", min_version="2.4", request_data_model=StatusReportRequest)
+@endpoint(
+    "workers.status_report", min_version="2.4", request_data_model=StatusReportRequest
+)
 def status_report(call: APICall, company_id, request: StatusReportRequest):
     worker_bll.status_report(
         company_id=company_id,