Feature/cursor paging

edfi_api_client/edfi_endpoint.py

@@ -11,6 +11,9 @@
 if TYPE_CHECKING:
     from edfi_api_client.edfi_client import EdFiClient
 
+from joblib import Parallel, delayed 
+from functools import partial
+
 
 class EdFiEndpoint:
     """
@@ -123,30 +126,30 @@ def ping(self, *, params: Optional[dict] = None, **kwargs) -> requests.Response:
         return res
 
 
-    def get(self, limit: Optional[int] = None, *, params: Optional[dict] = None, **kwargs) -> List[dict]:
+    def get(self, url: Optional[str] = None, limit: Optional[int] = None, *, params: Optional[dict] = None, **kwargs) -> List[dict]:
         """
         This method returns the rows from a single GET request using the exact params passed by the user.
 
         :return:
         """
-        logging.info(f"[Get {self.component}] Endpoint: {self.url}")
+        end_point = url or self.url
+        logging.info(f"[Get {self.component}] Endpoint: {end_point}")
 
         # Override init params if passed
         params = EdFiParams(params or self.params).copy()
         if limit:  # Override limit if passed
             params['limit'] = limit
 
         logging.info(f"[Get {self.component}] Parameters: {params}")
-        return self.client.session.get_response(self.url, params=params, **kwargs).json()
+
+        resp = self.client.session.get_response(end_point, params=params, **kwargs).json()
+        return resp
 
 
     def get_rows(self,
         *,
         params: Optional[dict] = None,  # Optional alternative params
         page_size: int = 100,
-        reverse_paging: bool = True,
-        step_change_version: bool = False,
-        change_version_step_size: int = 50000,
         **kwargs
     ) -> Iterator[dict]:
         """
@@ -163,20 +166,22 @@ def get_rows(self,
         :param max_wait:
         :return:
         """
-        paged_result_iter = self.get_pages(
+
+        paged_result_iter = self.get_pages_cursor(
             params=params,
-            page_size=page_size, reverse_paging=reverse_paging,
-            step_change_version=step_change_version, change_version_step_size=change_version_step_size,
+            page_size=page_size,
             **kwargs
         )
 
         for paged_result in paged_result_iter:
             yield from paged_result
 
 
-    def get_pages(self,
+    def get_pages_offset(self,
         *,
+        url: Optional[str] = None,
         params: Optional[dict] = None,  # Optional alternative params
+        limit: Optional[int] = None,
         page_size: int = 100,
         reverse_paging: bool = True,
         step_change_version: bool = False,
@@ -199,6 +204,11 @@ def get_pages(self,
         """
         # Override init params if passed
         paged_params = EdFiParams(params or self.params).copy()
+        end_point = url or self.url
+        logging.info(f"[Get {self.component}] Endpoint: {end_point}")
+
+        if limit:  # Override limit if passed
+            paged_params['limit'] = limit
 
         ### Prepare pagination variables, depending on type of pagination being used
         if step_change_version and reverse_paging:
@@ -210,19 +220,20 @@ def get_pages(self,
         elif step_change_version:
             logging.info(f"[Paged Get {self.component}] Pagination Method: Change Version Stepping")
             paged_params.init_page_by_offset(page_size)
-            paged_params.init_page_by_change_version_step(change_version_step_size)
-
+            paged_params.init_page_by_change_version_step(change_version_step_size)         
         else:
             logging.info(f"[Paged Get {self.component}] Pagination Method: Offset Pagination")
             paged_params.init_page_by_offset(page_size)
 
+        total_count = 0
         # Begin pagination-loop
         while True:
+            logging.info(f"[Get {self.component}] Parameters: {paged_params}")
 
             ### GET from the API and yield the resulting JSON payload
-            paged_rows = self.get(params=paged_params, **kwargs)
-            logging.info(f"[Get {self.component}] Retrieved {len(paged_rows)} rows.")
+            paged_rows = self.client.session.get_response(end_point, params=paged_params, **kwargs).json()
             yield paged_rows
+            logging.info(f"[Get {self.component}] Retrieved {len(paged_rows)} rows.")
 
             ### Paginate, depending on the method specified in arguments
             # Reverse offset pagination is only applicable during change-version stepping.
@@ -239,7 +250,7 @@ def get_pages(self,
                     except StopIteration:
                         logging.info(f"[Paged Get {self.component}] @ Change version exceeded max. Ending pagination.")
                         break
-
+                
             else:
                 # If no rows are returned, end pagination.
                 if len(paged_rows) == 0:
@@ -259,6 +270,67 @@ def get_pages(self,
                 else:
                     logging.info(f"@ Paginating offset...")
                     paged_params.page_by_offset()
+
+    def get_pages_cursor(self,
+        *,
+        url: Optional[str] = None,
+        params: Optional[dict] = None,  # Optional alternative params
+        limit: Optional[int] = None,
+        page_size: int = 100,
+        **kwargs
+    ) -> Iterator[List[dict]]:
+
+        # Override init params if passed
+        paged_params = EdFiParams(params or self.params).copy()
+        end_point = url or self.url
+        logging.info(f"[Get {self.component}] Endpoint: {end_point}")
+
+        if limit:  # Override limit if passed
+            paged_params['limit'] = limit
+
+        # Fall back to reverse-offset paging if incompatible with cursor paging
+        def _fall_back_to_pages_by_offset():
+            return self.get_pages_offset(
+                url = url,
+                params = params, 
+                limit = limit, 
+                page_size=page_size,
+                **kwargs
+            )
+
+        # Check ODS version compatibility for cursor paging
+        ods_version = tuple(map(int, self.client.get_ods_version().split(".")[:2]))
+        if ods_version < (7,3):
+            logging.warning(f"ODS {self.client.get_ods_version()} is incompatible. Cursor Paging requires v.7.3 or higher. Falling back to another paging method")
+            yield from _fall_back_to_pages_by_offset()
+            return
+        # deletes/key_changes cannot be retrieved with cursor paging
+        if self.get_deletes or self.get_key_changes:
+            logging.warning(f"Cursor Paging does not support deletes/key_changes. Falling back to another paging method")
+            yield from _fall_back_to_pages_by_offset()
+            return
+
+        logging.info(f"[Paged Get {self.component}] Pagination Method: Cursor Paging")
+
+        ###  Prepare pagination variables 
+        ###  First request should not have any `page_token` and `page_size` defined
+        paged_params.init_page_by_token(page_token = None, page_size = None)            
+
+        # Begin pagination loop
+        while True:
+            logging.info(f"[Get {self.component}] Parameters: {paged_params}")
+
+            result = self.client.session.get_response(end_point, params = paged_params, **kwargs)
+            paged_rows = result.json()
+            logging.info(f"[Get {self.component}] Retrieved {len(paged_rows)} rows")
+            yield paged_rows
+
+            logging.info(f"[Paged Get {self.component}] @ Cursor paging ...")
+            if not result.headers.get("Next-Page-Token"):
+                logging.info(f"[Paged Get {self.component}] @ Retrieved zero rows. Ending pagination.")
+                break
+            paged_params.init_page_by_token(page_token = result.headers.get("Next-Page-Token"), page_size = page_size)
+
 
 
     def get_total_count(self, *, params: Optional[dict] = None, **kwargs) -> int:

edfi_api_client/edfi_params.py

@@ -25,6 +25,8 @@ def __init__(self,
         # These parameters are only used during pagination. They must be explicitly initialized.
         self.page_size = None
         self.change_version_step_size = None
+        self.page_token = None
+        self.number = None
 
 
     def copy(self) -> 'EdFiParams':
@@ -190,3 +192,28 @@ def reverse_page_by_offset(self):
 
         if self['offset'] < 0:
             raise StopIteration
+
+    def init_page_by_token(self, page_token: str, page_size: int):
+        """
+
+        :param page_size: 
+        :param page_token:
+        :return:
+        """
+
+        # Cursor paging behavior: page_token is required when page_size is specified.
+        # - If page_token is None: first request, do NOT include page_size
+        # - If page_token is present: include page_token and page_size
+        self.page_size = page_size
+        self.page_token = page_token
+
+        if page_token is None:
+            self.pop("pageToken", None)
+            self.pop("page_size", None)
+        else:
+            self["pageToken"] = self.page_token
+            self["page_size"] = self.page_size 
+
+
+
+