Respect ingestion limit in process_batch

[stacimc]

Description

I think this should fix the final big edge case with ingestion limit :)

Before this PR, when an INGESTION_LIMIT was set in a provider script, we would cap the batch_limit to that value, and then check if the limit was exceeded after each call to process_batch. This leaves an edge case when a provider doesn't pass batch_limit to their API, and may get back huge batches that are larger than our limit.

This PR adds one extra check for the ingestion_limit, inside process_batch itself. Now we will halt processing even within a batch when the limit is reached.

Testing Instructions

Metropolitan is a good example of the buggy behavior. You can observe the bug in action by setting an INGESTION_LIMIT Airflow variable to 100 locally and running metropolitan_workflow. You'll see that it appears to ignore the limit and continues processing after 100.

Then checkout this branch and try again. You could also try verifying that the limit is still respected on other providers like Cleveland.

Checklist

• My pull request has a descriptive title (not a vague title like Update index.md).
• My pull request targets the default branch of the repository (main) or a parent feature branch.
• My commit messages follow best practices.
• My code follows the established code style of the repository.
• I added or updated tests for the changes I made (if applicable).
• I added or updated documentation (if applicable).
• I tried running the project locally and verified that there are no visible errors.

Developer Certificate of Origin

Developer Certificate of Origin

Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
1 Letterman Drive
Suite D4700
San Francisco, CA, 94129

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.


Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.

[AetherUnbound]

Nice catch! Do you think we'd be able to remove any of the higher level checks if we have this one at the lower level? 😮

[stacimc]

Do you think we'd be able to remove any of the higher level checks if we have this one at the lower level? 😮

Weirdly, I don't think so! Or rather, we could, but while it would stop the ingester from committing any new records past the limit, it wouldn't stop it from trying to fetch additional batches (it would keep fetching and return early each time). The current checks are all needed to make sure ingestion totally halts immediately.

There's two other checks:

• [1] At the beginning of ingest_records. We need this one to stop the ingester from attempting to start ingestion again in providers that override ingest_records to be called multiple times (like Finnish).
• [2] At the end of ingest_records -- I thought we might be able to get rid of this one, but we need it to set should_continue to False and stop the loop.

It is surprisingly complicated to implement in a way that suits all our providers!

[obulat]

LGTM!

[krysal]

Works great! Also, thanks for the context, I was feeling a dejavú like, some work in this area has been done before.

[stacimc]

I was feeling a dejavú like, some work in this area has been done before.

You and me both! 😄