CASMCMS-9225: Added basic paging ability for GET requests to list com…

…ponents (#396) * CASMCMS-9225: Added basic paging ability for GET requests to list components * Apply suggestions from code review * Update src/bos/server/redis_db_utils.py * Combine id_list and tenant_components in filter function; Immediately return empty list if id_list&tenant_components is empty * Fix mistake in code review update * Fix bug in new components logic * Apply suggestions from code review Co-authored-by: Jason Sollom <[email protected]> --------- Co-authored-by: Jason Sollom <[email protected]>
Cray-HPE · Dec 17, 2024 · cea748f · cea748f
1 parent a529751
commit cea748f
Show file tree

Hide file tree

Showing 4 changed files with 134 additions and 54 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Added
+- Added basic paging ability for `GET` requests for `components`.
 
 ## [2.31.0] - 2024-11-01
 ### Removed

diff --git a/api/openapi.yaml.in b/api/openapi.yaml.in
@@ -1691,6 +1691,21 @@ paths:
           in: query
           description: |-
             Retrieve the Components with the given status.
+        - name: start_after_id
+          schema:
+            $ref: '#/components/schemas/V2ComponentId'
+          in: query
+          description: |-
+            Begin listing Components after the specified ID. Used for paging.
+        - name: page_size
+          schema:
+            type: integer
+            minimum: 0
+            maximum: 1048576
+          in: query
+          description: |-
+            Maximum number of Components to include in response. Used for paging. 0 means no limit
+            (which is the same as not specifying this parameter).
       description: |-
         Retrieve the full collection of Components in the form of a
         ComponentArray. Full results can also be filtered by query

diff --git a/src/bos/server/controllers/v2/components.py b/src/bos/server/controllers/v2/components.py
@@ -21,6 +21,7 @@
 # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 # OTHER DEALINGS IN THE SOFTWARE.
 #
+from functools import partial
 import logging
 
 import connexion
@@ -46,7 +47,9 @@ def get_v2_components(ids="",
                       session=None,
                       staged_session=None,
                       phase=None,
-                      status=None):
+                      status=None,
+                      start_after_id=None,
+                      page_size=0):
     """Used by the GET /components API operation
 
     Allows filtering using a comma separated list of ids.
@@ -55,7 +58,6 @@ def get_v2_components(ids="",
         "GET /v2/components invoked get_v2_components with ids=%s enabled=%s session=%s "
         "staged_session=%s phase=%s status=%s", ids, enabled, session,
         staged_session, phase, status)
-    id_list = []
     if ids:
         try:
             id_list = ids.split(',')
@@ -64,21 +66,24 @@ def get_v2_components(ids="",
             return connexion.problem(status=400,
                                      title="Error parsing the ids provided.",
                                      detail=str(err))
+    else:
+        id_list = None
     tenant = get_tenant_from_header()
     LOGGER.debug("GET /v2/components for tenant=%s with %d IDs specified",
-                 tenant, len(id_list))
+                 tenant, len(id_list) if id_list else 0)
     response = get_v2_components_data(id_list=id_list,
                                       enabled=enabled,
                                       session=session,
                                       staged_session=staged_session,
                                       phase=phase,
                                       status=status,
-                                      tenant=tenant)
+                                      tenant=tenant,
+                                      start_after_id=start_after_id,
+                                      page_size=page_size,
+                                      delete_timestamp=True)
     LOGGER.debug(
         "GET /v2/components returning data for tenant=%s on %d components",
         tenant, len(response))
-    for component in response:
-        del_timestamp(component)
     return response, 200
 
 
@@ -88,47 +93,87 @@ def get_v2_components_data(id_list=None,
                            staged_session=None,
                            phase=None,
                            status=None,
-                           tenant=None):
+                           tenant=None,
+                           start_after_id=None,
+                           page_size=0,
+                           delete_timestamp: bool = False):
     """Used by the GET /components API operation
 
     Allows filtering using a comma separated list of ids.
     """
-    response = []
-    if id_list:
-        for component_id in id_list:
-            data = DB.get(component_id)
-            if data:
-                response.append(data)
+    tenant_components = None if tenant is None else get_tenant_component_set(tenant)
+
+    if id_list is not None:
+        id_set = set(id_list)
+        if tenant_components is not None:
+            id_set.intersection_update(tenant_components)
     else:
-        # TODO: On large scale systems, this response may be too large
-        # and require paging to be implemented
-        response = DB.get_all()
-    # The status must be set before using _matches_filter as the status is one of the
-    # matching criteria.
-    response = [_set_status(r) for r in response if r]
-    if enabled is not None or session is not None or staged_session is not None or \
-       phase is not None or status is not None:
-        response = [
-            r for r in response if _matches_filter(
-                r, enabled, session, staged_session, phase, status)
-        ]
-    if tenant:
-        tenant_components = get_tenant_component_set(tenant)
-        limited_response = [
-            component for component in response
-            if component["id"] in tenant_components
-        ]
-        response = limited_response
-    return response
+         id_set = tenant_components
+
+    # If id_set is not None but is empty, that means no components in the system
+    # will match our filter, so we can return an empty list immediately.
+    if id_set is not None and not id_set:
+        return []
+
+    if any([id_set, enabled, session, staged_session, phase, status]):
+        _component_filter_func = partial(_filter_component,
+                                         id_set=id_set,
+                                         enabled=enabled,
+                                         session=session or None,
+                                         staged_session=staged_session or None,
+                                         phase=phase or None,
+                                         status=status or None)
+    elif delete_timestamp:
+        _component_filter_func = partial(_set_status,
+                                         delete_timestamp=delete_timestamp)
+    else:
+        _component_filter_func = _set_status
+
+    return DB.get_all_filtered(filter_func=_component_filter_func,
+                               start_after_key=start_after_id,
+                               page_size=page_size)
+
+
+def _filter_component(data: dict,
+                      id_set=None,
+                      enabled=None,
+                      session=None,
+                      staged_session=None,
+                      phase=None,
+                      status=None,
+                      delete_timestamp: bool = False) -> dict | None:
+    # Do all of the checks we can before calculating status, to avoid doing it needlessly
+    if id_set is not None and data["id"] not in id_set:
+        return None
+    if enabled is not None and data.get('enabled', None) != enabled:
+        return None
+    if session is not None and data.get('session', None) != session:
+        return None
+    if staged_session is not None and \
+       data.get('staged_state', {}).get('session', None) != staged_session:
+        return None
+    updated_data = _set_status(data)
+
+    status_data = updated_data.get('status')
+    if phase is not None and status_data.get('phase') != phase:
+        return None
+    if status is not None and status_data.get('status') not in status.split(
+            ','):
+        return None
+    if delete_timestamp:
+        del_timestamp(updated_data)
+    return updated_data
 
 
-def _set_status(data):
+def _set_status(data, delete_timestamp: bool = False):
     """
     This sets the status field of the overall status.
     """
     if "status" not in data:
         data["status"] = {"phase": "", "status_override": ""}
     data['status']['status'] = _calculate_status(data)
+    if delete_timestamp:
+        del_timestamp(data)
     return data
 
 
@@ -147,12 +192,13 @@ def _calculate_status(data):
 
     phase = status_data.get('phase', '')
     component = data.get('id', '')
-    last_action = data.get('last_action', {}).get('action', '')
+    last_action_dict = data.get('last_action', {})
+    last_action = last_action_dict.get('action', '')
 
     status = status = Status.stable
     if phase == Phase.powering_on:
-        if last_action == Action.power_on and not data.get(
-                'last_action', {}).get('failed', False):
+        if last_action == Action.power_on and not last_action_dict.get(
+                'failed', False):
             status = Status.power_on_called
         else:
             status = Status.power_on_pending
@@ -171,23 +217,6 @@ def _calculate_status(data):
     return status
 
 
-def _matches_filter(data, enabled, session, staged_session, phase, status):
-    if enabled is not None and data.get('enabled', None) != enabled:
-        return False
-    if session is not None and data.get('session', None) != session:
-        return False
-    if staged_session is not None and \
-       data.get('staged_state', {}).get('session', None) != staged_session:
-        return False
-    status_data = data.get('status')
-    if phase is not None and status_data.get('phase') != phase:
-        return False
-    if status is not None and status_data.get('status') not in status.split(
-            ','):
-        return False
-    return True
-
-
 @dbutils.redis_error_handler
 def put_v2_components():
     """Used by the PUT /components API operation"""

diff --git a/src/bos/server/redis_db_utils.py b/src/bos/server/redis_db_utils.py
@@ -21,9 +21,11 @@
 # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 # OTHER DEALINGS IN THE SOFTWARE.
 #
+from itertools import batched
 import functools
 import json
 import logging
+from typing import Callable, Optional
 
 import connexion
 import redis
@@ -106,6 +108,38 @@ def get_all(self):
             data.append(single_data)
         return data
 
+    def get_all_filtered(self,
+                         filter_func: Callable[[dict], dict | None],
+                         start_after_key: Optional[str] = None,
+                         page_size: int = 0) -> list[dict]:
+        """
+        Get an array of data for all keys after passing them through the specified filter
+        (discarding any for which the filter returns None)
+        If start_after_id is specified, all ids lexically <= that id will be skipped.
+        If page_size is specified, the number of items in the returned list will be equal to or less than the page_size.
+        More elements may remain and additional queries will be needed to acquire them.
+        """
+        data = []
+        for value in self.iter_values(start_after_key):
+            filtered_value = filter_func(value)
+            if filtered_value is not None:
+                data.append(filtered_value)
+                if page_size and len(data) == page_size:
+                    break
+        return data
+
+    def iter_values(self, start_after_key: Optional[str] = None):
+        """
+        Iterate through every item in the database. Parse each item as JSON and yield it.
+        If start_after_key is specified, skip any keys that are lexically <= the specified key.
+        """
+        all_keys = sorted({k.decode() for k in self.client.scan_iter()})
+        if start_after_key is not None:
+            all_keys = [k for k in all_keys if k > start_after_key]
+        for next_keys in batched(all_keys, 500):
+            for datastr in self.client.mget(next_keys):
+                yield json.loads(datastr) if datastr else None
+
     def get_all_as_dict(self):
         """Return a mapping from all keys to their corresponding data
            Based on https://github.com/redis/redis-py/issues/984#issuecomment-391404875