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,29 +126,34 @@ 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) -> requests.Response:
         """
         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)
+        return resp
 
 
     def get_rows(self,
         *,
         params: Optional[dict] = None,  # Optional alternative params
         page_size: int = 100,
-        reverse_paging: bool = True,
+        reverse_paging: bool = False,
         step_change_version: bool = False,
+        cursor_paging: bool = True,
+        partitioning: bool = False,
         change_version_step_size: int = 50000,
         **kwargs
     ) -> Iterator[dict]:
@@ -165,7 +173,7 @@ def get_rows(self,
         """
         paged_result_iter = self.get_pages(
             params=params,
-            page_size=page_size, reverse_paging=reverse_paging,
+            page_size=page_size, reverse_paging=reverse_paging, cursor_paging= cursor_paging, partitioning=partitioning,
             step_change_version=step_change_version, change_version_step_size=change_version_step_size,
             **kwargs
         )
@@ -180,7 +188,10 @@ def get_pages(self,
         page_size: int = 100,
         reverse_paging: bool = True,
         step_change_version: bool = False,
+        cursor_paging: bool = False,
+        partitioning: bool = False,
         change_version_step_size: int = 50000,
+        number: Optional[int] = None,
         **kwargs
     ) -> Iterator[List[dict]]:
         """
@@ -211,18 +222,44 @@ def get_pages(self,
             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)
-
+
+        elif cursor_paging and partitioning:
+            ods_version = tuple(map(int, self.client.get_ods_version().split(".")[:2]))
+            if ods_version < (7,3):
+                raise ValueError(f"ODS {self.client.get_ods_version()} is incompatible. Cursor Paging with Partitions requires v.7.3 or higher. Ending pagination")
+            if self.get_deletes or self.get_key_changes:
+                raise ValueError(f"Cursor Paging with Partitions does not support deletes/key_changes. Ending pagination")
+
+            token_url = f"{self.url.rstrip('/')}/partitions"
+            paged_tokens = self.get(url = token_url, params = paged_params.init_page_by_partitions(number = number) , **kwargs).json().get("pageTokens")
+            logging.info(f"[Paged Get {self.component}] Pagination Method: Cursor Paging with Partitions")
+            logging.info(f"[Get {self.component}] Retrieved {len(paged_tokens)} token(s): {paged_tokens}")
+
+        elif cursor_paging:
+            ods_version = tuple(map(int, self.client.get_ods_version().split(".")[:2]))
+            if ods_version < (7,3):
+                raise ValueError(f"ODS {self.client.get_ods_version()} is incompatible. Cursor Paging requires v.7.3 or higher. Ending pagination")
+            if self.get_deletes or self.get_key_changes:
+                raise ValueError(f"Cursor Paging does not support deletes/key_changes. Ending pagination")
+            logging.info(f"ODS {self.client.get_ods_version()}")
+            logging.info(f"[Paged Get {self.component}] Pagination Method: Cursor Paging")
+            # First request should not have any `page_token` defined
+            paged_params.init_page_by_token(page_token = None, page_size = None)            
         else:
             logging.info(f"[Paged Get {self.component}] Pagination Method: Offset Pagination")
             paged_params.init_page_by_offset(page_size)
 
+
         # Begin pagination-loop
         while True:
 
             ### 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.")
-            yield paged_rows
+            if not partitioning:
+                result = self.get(params=paged_params, **kwargs)
+                paged_rows = result.json()
+                logging.info(f"[Get {self.component}] Retrieved {len(paged_rows)} rows.")
+
+                yield paged_rows
 
             ### Paginate, depending on the method specified in arguments
             # Reverse offset pagination is only applicable during change-version stepping.
@@ -239,7 +276,35 @@ def get_pages(self,
                     except StopIteration:
                         logging.info(f"[Paged Get {self.component}] @ Change version exceeded max. Ending pagination.")
                         break
-
+            elif cursor_paging and partitioning:
+                def partitioning_with_token(token): 
+                    results = [] 
+                    p = paged_params.copy() 
+                    while True: 
+                        p.init_page_by_token(page_token=token, page_size=page_size) 
+                        result = self.get(params=p, **kwargs) 
+                        paged_rows = result.json() 
+                        results.extend(paged_rows) 
+                        if len(paged_rows) == 0: 
+                            logging.info(f"[Paged Get {self.component}] @ Retrieved zero rows for token: {token}. Ending pagination.") 
+                            break 
+                        logging.info(f"[Get {self.component}] Retrieved {len(paged_rows)} rows for token: {token}.")                     
+                        token = result.headers.get("Next-Page-Token")
+                    return results
+
+                results = [element for sublist in Parallel(n_jobs=len(paged_tokens), backend="threading")(delayed(partitioning_with_token)(token) for token in paged_tokens) for element in sublist]
+                logging.info(f"{len(results)}")
+                for paged_page in results:
+                    yield paged_page
+                return
+
+            elif cursor_paging : 
+                logging.info(f"[Paged Get {self.component}] @ Cursor paging ...")
+                paged_params.init_page_by_token( page_token = result.headers.get("Next-Page-Token"), page_size = page_size)
+                if not result.headers.get("Next-Page-Token"):
+                    logging.info(f"[Paged Get {self.component}] @ Retrieved zero rows. Ending pagination.")
+                    break
+
             else:
                 # If no rows are returned, end pagination.
                 if len(paged_rows) == 0:
@@ -259,7 +324,7 @@ def get_pages(self,
                 else:
                     logging.info(f"@ Paginating offset...")
                     paged_params.page_by_offset()
-
+            
 
     def get_total_count(self, *, params: Optional[dict] = None, **kwargs) -> int:
         """
@@ -319,6 +384,18 @@ def _get_attributes_from_swagger(self):
             'has_deletes': (self.namespace, self.name) in swagger.deletes,
         }
 
+    def partitioning_with_token(self, params: Optional[dict] = None, token: str = None, page_size: str = None, **kwargs ):
+
+        params.init_page_by_token(page_token=token, page_size=page_size)
+        result = params.get(params=params, **kwargs)
+        paged_rows = result.json()
+        yield paged_rows
+        logging.info(f"[Get {self.component}] Retrieved {len(paged_rows)} rows.")
+
+        token = result.headers.get("Next-Page-Token")
+        if not token:
+            logging.info(f"[Paged Get {self.component}] @ Retrieved zero rows. Ending pagination.")
+            return False
 
 class EdFiResource(EdFiEndpoint):
     component: str = 'Resource'

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,36 @@ 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("pageSize", None)
+        else:
+            self["page_token"] = self.page_token
+            self["page_size"] = self.page_size 
+
+    def init_page_by_partitions(self, number: int):
+        """
+
+        :param number: 
+        :return:
+        """
+
+        self.number = number
+        self["number"] = number
+
+