diff --git a/docs/rest/ProjectsApi.md b/docs/rest/ProjectsApi.md index e8cd5e4d..77ede151 100644 --- a/docs/rest/ProjectsApi.md +++ b/docs/rest/ProjectsApi.md @@ -1547,7 +1547,7 @@ with gretel_client.rest.ApiClient(configuration) as api_client: api_instance = projects_api.ProjectsApi(api_client) project_id = "project_id_example" # str | Project id model_id = "model_id_example" # str | Model id - status = "completed" # str | + status = "status_example" # str | skip = 1 # int | The number of records being skipped before returning the next set. (optional) limit = 1 # int | The number of records returned in each result set. (optional) expand = [ diff --git a/src/gretel_client/projects/models.py b/src/gretel_client/projects/models.py index 7367325a..a642838d 100644 --- a/src/gretel_client/projects/models.py +++ b/src/gretel_client/projects/models.py @@ -434,26 +434,25 @@ def _do_cancel_job(self): def get_record_handlers(self) -> Iterator[RecordHandler]: """Returns a list of record handlers associated with the model.""" - for status in Status: - # Setting the start value for offset in each query, limiting each result set to 10 records - offset = 0 - limit = 10 - while True: - handlers = ( - self._projects_api.query_record_handlers( - project_id=self.project.project_id, - model_id=self.model_id, - status=status.value, - skip=offset, - limit=limit, - ) - .get("data") - .get("handlers") + # Setting the start value for offset in each query, limiting each result set to 10 records + offset = 0 + limit = 10 + while True: + handlers = ( + self._projects_api.query_record_handlers( + project_id=self.project.project_id, + model_id=self.model_id, + skip=offset, + status=",".join(status.value for status in Status), + limit=limit, ) - # will break once no more handlers remain in the result set in the last pagination - if len(handlers) > 0: - for handler in handlers: - yield RecordHandler(self, record_id=handler["uid"]) - offset += limit - else: - break + .get("data") + .get("handlers") + ) + # will break once no more handlers remain in the result set in the last pagination + if len(handlers) > 0: + for handler in handlers: + yield RecordHandler(self, record_id=handler["uid"]) + offset += limit + else: + break diff --git a/src/gretel_client/rest/api/projects_api.py b/src/gretel_client/rest/api/projects_api.py index ca340436..d1b5285f 100644 --- a/src/gretel_client/rest/api/projects_api.py +++ b/src/gretel_client/rest/api/projects_api.py @@ -2359,7 +2359,6 @@ def __query_record_handlers(self, project_id, model_id, status, **kwargs): ], "nullable": [], "enum": [ - "status", "expand", ], "validation": [], @@ -2367,15 +2366,6 @@ def __query_record_handlers(self, project_id, model_id, status, **kwargs): root_map={ "validations": {}, "allowed_values": { - ("status",): { - "COMPLETED": "completed", - "ERROR": "error", - "PENDING": "pending", - "ACTIVE": "active", - "LOST": "lost", - "CREATED": "created", - "CANCELLED": "cancelled", - }, ("expand",): {"ARTIFACTS": "artifacts"}, }, "openapi_types": { diff --git a/tests/gretel_client/test_models.py b/tests/gretel_client/test_models.py index 0b37f50d..ca41633e 100644 --- a/tests/gretel_client/test_models.py +++ b/tests/gretel_client/test_models.py @@ -301,11 +301,7 @@ def test_goes_through_records( m.submit(runner_mode=RunnerMode.LOCAL) def mock_record_handlers(status, skip, limit, *args, **kwargs): - handlers = ( - [{"uid": 0} for _ in range(min(limit, num_records - skip))] - if status == "completed" - else [] - ) + handlers = [{"uid": 0} for _ in range(min(limit, num_records - skip))] return {"data": {"handlers": handlers}} m._projects_api.query_record_handlers.side_effect = mock_record_handlers