Allow botlist API call to respond to quarantined: and is_dead:

appengine/swarming/handlers_endpoints.py

Index: appengine/swarming/handlers_endpoints.py
diff --git a/appengine/swarming/handlers_endpoints.py b/appengine/swarming/handlers_endpoints.py
index 5ced064a782d16430703831503df9e22a64268ce..766c374599391475c5467c0731040ee66a207ae1 100644
--- a/appengine/swarming/handlers_endpoints.py
+++ b/appengine/swarming/handlers_endpoints.py
@@ -608,6 +608,30 @@ class SwarmingBotService(remote.Service):
 @swarming_api.api_class(resource_name='bots', path='bots')
 class SwarmingBotsService(remote.Service):
   """Bots-related API."""
+
+  def _filter_dimensions(self, q, request):

[M-A Ruel, 2016/08/09 18:05:27]

maybe worth passing request.dimensions as an argument instead of request, so
it's cleaner here.

[kjlubick, 2016/08/09 19:25:03]

Done.

+    """Filters a ndb.Query for BotInfo based on dimensions in the request."""
+    for d in request.dimensions:
+      parts = d.split(':', 1)
+      if len(parts) != 2 or any(i.strip() != i or not i for i in parts):
+        raise endpoints.BadRequestException('Invalid dimensions')
+      q = q.filter(bot_management.BotInfo.dimensions_flat == d)
+    return q
+
+  def _filter_availability(self, q, request, now):

[M-A Ruel, 2016/08/09 18:05:27]

same here
def _filter_availability(self, q, is_busy, quarantined, is_dead, now):

Actually, these functions should them be moved to server/bot_management.py. This
way, if we need to refactor the way the entity BotInfo is stored to optimize the
indexes, most of the changes are self-contained in bot_management.py.

From bot_management.py, you should raise ValueError(). You can use a
try/except+raise to convert to endpoints.BadRequestException().

[kjlubick, 2016/08/09 19:25:03]

Done.

+    """Filters a ndb.Query for BotInfo based on quarantined/is_dead."""
+    val = swarming_rpcs.to_bool(request.quarantined)
+    if val is not None:
+      q = q.filter(bot_management.BotInfo.quarantined == val)
+
+    dt = datetime.timedelta(seconds=config.settings().bot_death_timeout_secs)
+    timeout = now - dt
+    if request.is_dead == swarming_rpcs.ThreeStateBool.TRUE:
+      q = q.filter(bot_management.BotInfo.last_seen_ts < timeout)
+    elif request.is_dead == swarming_rpcs.ThreeStateBool.FALSE:
+      q = q.filter(bot_management.BotInfo.last_seen_ts > timeout)
+    return q
+
   @gae_ts_mon.instrument_endpoint()
   @auth.endpoints_method(
       swarming_rpcs.BotsRequest, swarming_rpcs.BotList,
@@ -620,14 +644,10 @@ class SwarmingBotsService(remote.Service):
     """
     logging.info('%s', request)
     now = utils.utcnow()
-    q = bot_management.BotInfo.query().order(bot_management.BotInfo.key)
-    for d in request.dimensions:
-      if not ':' in d:
-        raise endpoints.BadRequestException('Invalid dimensions')
-      parts = d.split(':', 1)
-      if len(parts) != 2 or any(i.strip() != i or not i for i in parts):
-        raise endpoints.BadRequestException('Invalid dimensions')
-      q = q.filter(bot_management.BotInfo.dimensions_flat == d)
+    q = bot_management.BotInfo.query()
+    q = self._filter_dimensions(q, request)
+    q = self._filter_availability(q, request, now)
+
     bots, cursor = datastore_utils.fetch_page(q, request.limit, request.cursor)
     return swarming_rpcs.BotList(
         cursor=cursor,
@@ -645,25 +665,50 @@ class SwarmingBotsService(remote.Service):
     logging.info('%s', request)
     now = utils.utcnow()
     q = bot_management.BotInfo.query()

[M-A Ruel, 2016/08/09 18:05:26]

you can probably inline lines 667 and 668, and same for 647 and 648.

[kjlubick, 2016/08/09 19:25:03]

Cannot, it is too long

-    for d in request.dimensions:
-      parts = d.split(':', 1)
-      if len(parts) != 2 or any(i.strip() != i or not i for i in parts):
-        raise endpoints.BadRequestException('Invalid dimensions: %s' % d)
-      q = q.filter(bot_management.BotInfo.dimensions_flat == d)
-    f_count = q.count_async()
+    q = self._filter_dimensions(q, request)
+
+    quarantined = swarming_rpcs.to_bool(request.quarantined)
+    is_dead = swarming_rpcs.to_bool(request.is_dead)
     dt = datetime.timedelta(seconds=config.settings().bot_death_timeout_secs)
     timeout = now - dt
-    f_dead = q.filter(
-        bot_management.BotInfo.last_seen_ts < timeout).count_async()
-    f_quarantined = q.filter(
-        bot_management.BotInfo.quarantined == True).count_async()
-    f_busy = q.filter(bot_management.BotInfo.is_busy == True).count_async()
-    return swarming_rpcs.BotsCount(
-        count=f_count.get_result(),
-        quarantined=f_quarantined.get_result(),
-        dead=f_dead.get_result(),
-        busy=f_busy.get_result(),
-        now=now)
+
+    if quarantined is None and is_dead is None:
+      f_count = q.count_async()
+      f_dead = q.filter(

[M-A Ruel, 2016/08/09 18:05:26]

you can do:
f_dead = self._filter_availability(q, True, False).count_async()

[kjlubick, 2016/08/09 19:25:03]

Done.

+          bot_management.BotInfo.last_seen_ts < timeout).count_async()
+      f_quarantined = q.filter(

[M-A Ruel, 2016/08/09 18:05:27]

same here

[kjlubick, 2016/08/09 19:25:03]

Done.

+          bot_management.BotInfo.quarantined == True).count_async()
+      f_busy = q.filter(bot_management.BotInfo.is_busy == True).count_async()
+      return swarming_rpcs.BotsCount(
+          count=f_count.get_result(),
+          quarantined=f_quarantined.get_result(),
+          dead=f_dead.get_result(),
+          busy=f_busy.get_result(),
+          now=now)
+    elif is_dead is not None:
+      if is_dead:
+        q = q.filter(bot_management.BotInfo.last_seen_ts < timeout)
+      else:
+        q = q.filter(bot_management.BotInfo.last_seen_ts > timeout)
+      count = q.count()
+      # Dead bots, by definition, are not quarantined, nor are they busy.
+      return swarming_rpcs.BotsCount(
+          count=count,
+          quarantined=0,
+          dead=count,
+          busy=0,
+          now=now)
+    else:
+      # Quarantined bots, by definition, are not dead, nor are they busy.
+      q = q.filter(bot_management.BotInfo.quarantined == quarantined)
+      count = q.count()
+      return swarming_rpcs.BotsCount(
+          count=count,
+          quarantined=count,
+          dead=0,
+          busy=0,
+          now=now)
+
 
   @gae_ts_mon.instrument_endpoint()
   @auth.endpoints_method(