Support partial parsing

[dwreeves]

Description

dbt uses partial_parse.msgpack to make rendering things a lot faster. This PR adds support for partial_parse.msgpack in the following places:

• ExecutionMode.LOCAL
• ExecutionMode.VIRTUALENV
• LoadMode.DBT_LS

This PR also allows users to explicitly turn off partial parsing. If this is done, then --no-partial-parse will be passed as an explicit flag in all dbt command invocations (i.e. all ExecutionModes and LoadMode.DBT_LS, albeit not the dbt deps invocation.)

This should address some performance complaints that users have, e.g. this message from Slack: https://apache-airflow.slack.com/archives/C059CC42E9W/p1704483361206829 Albeit, this user will also need to provide a partial_parse.msgpack.

My experimentation and looking at dbt-core source code confirms that dbt does not use manifest.json when partial parsing. It appears that these are just output artifacts, but not input artifacts. Only partial_parse.msgpack is used. (There are a couple ways to confirm this other than just checking source code

Also, I added a minor refactor of a feature I added a year ago: I added send_sigint() to the custom subprocess hook, since this custom subprocess hook did not exist back when I added it (if you want me to split this refactor into a different PR then let me know).

API change

I decided the best way to go about this would be to just add a partial_parse: bool to both the execution config and render config. For example:

dbt_group = DbtTaskGroup(
    ...,
    execution_config=ExecutionConfig(
        ...,
        partial_parse=True
    ),
    render_config=RenderConfig(
        ...,
        partial_parse=False
    )
)

That said, in all honesty users will not need to set this at all, except unless they want to suppress the little warning message about not being able to find the partial_parse.msgpack. This is because by default this addition searches for a msgpack if it exists, which is already the existing behavior in a sense, except right now the msgpack file never exists (dbt does look for it though).

When inserting into the AbstractDbtBaseOperator, I did not use global_boolean_flags. See the subsection below for why.

Other execution performance improvements

The main reason I am adding this feature is that it should dramatically improve performance for users. However, it is not the only way to improve

It's possible that we should add a way to add the flag --no-write-json as an explicit kwarg to the dbt base operator. Right now users can support this via dbt_cmd_global_flags=["--no-write-json"]. Some users (e.g. those using Elementary, or those using the dbt local operator callback kwarg) will want to write the JSON, but I suspect the majority of users will not. Similarly, --log-level-file is not used at all, and at minimum dbt should work best the vast majority of time with --no-cache-selected-only set.

It's also possible there should be some sort of "performant" mode that automatically sets all these defaults for optimal performance:

• --no-write-json
• --log-level-file=none
• --no-cache-selected-only

Perhaps a "performant" config would be too cumbersome to implement (I would agree with that). In which case the docs could also have a section on performance tips.

A note on global_boolean_flags

I did not add the partial parse support to global_boolean_flags because it doesn't quite fit into the existing paradigm for this. Right now the default for each of these global_boolean_flags is False, whereas the default for partial parse is actually True. This makes fitting it in awkward.

I think it's possible that just having a tuple[str] is insufficient here, as some flags you may want to add (not just --no-partial-parse but also --no-write-json are by default True and must be explicitly turned off. Meaning that the parsing Cosmos does with flags of '--{flag.replace("_", "-")}' is ineffective for flags like this.

Right now, we have an example of putting no in front with no_version_check. Meaning that the default behavior of version checking is True, but the flag itself starts as negated so the default is actually False.

My proposal is, instead of global_boolean_flags: tuple[str], this should instead be global_boolean_flags: tuple[str | tuple[str, str | None, str | None]]. In the case that a global flag is a tuple[str, str | None, str | None], then the first arg should be the flag, the second should be "if true," and the third should be "if false." None indicates, when true/false (respectively), then do nothing.

For example:

class AbstractDbtBaseOperator(BaseOperator, metaclass=ABCMeta):
    ...
    global_boolean_flags = (
        ("no_version_check", "--no-version-check", None),
        ("cache_selected_only", "-cache-selected-only", None),
        ("partial_parse", None, "--no-partial-parse"),
    )

And Cosmos want to support str parsing for backwards compatibility. It's pretty straightforward to convert the data type:

if isinstance(flag, str):
    flag = (flag, '--{flag.replace("_", "-")}', None)

Related Issue(s)

• Resolves DBT_MANIFEST mode does not use manifest for task execution #791
• Partially resolves Option to copy partial_parse.msgpack or target/ dir into tmpdir. #785
  • Option to copy partial_parse.msgpack or target/ dir into tmpdir. #785 should probably be split up into two different stages: (1) support for partial parsing (2) (a) dbt project dir / manifest / partial_parse.msgpack is allowed to come from cloud storage. (b) dbt compile is able to dump into cloud storage.
• Closes Tasks are running full parse each time  #722

Breaking Change?

Should not break anything. This doesn't do anything when partial_parse.msgpack is missing, and the default behavior (partial_parse=True) does not alter the dbt cmd flags.

Checklist

• I have made corresponding changes to the documentation (if required)
• I have added tests that prove my fix is effective or that my feature works

[netlify]

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Name	Link
🔨 Latest commit	a9e9412
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/65d3b14bceb262000716e169

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (9af7067) 94.72% compared to head (a9e9412) 95.07%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #800      +/-   ##
==========================================
+ Coverage   94.72%   95.07%   +0.35%     
==========================================
  Files          56       56              
  Lines        2520     2540      +20     
==========================================
+ Hits         2387     2415      +28     
+ Misses        133      125       -8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dwreeves]

I don't know how the coverage thing works and if it is checking for 100% coverage of new commits vs just making sure the coverage goes up, but I helped increase the coverage of the code by adding tests for the subprocess hook.

[jlaneve]

This is awesome for two reasons: the change in and of itself, and the writeup on where we should look in the future. Apologies it's taken longer to review, but I'm going to check this out today.

I agree the performant mode is nice in theory but may be difficult to maintain. Maybe we should create a docs page specifically on performance with our recommendations? Inclusive of this flag, using the manifest parsing method, install LibYAML, etc. That way, for users who care about performance, there's a single doc for all optimization recommendations.

Quick question re: how this partial parsing plays out in practice. My assumption is that a user would:

• run dbt parse locally / in CICD (pre-deploying), which generates the necessary artifacts
• push the entire project to the Airflow deployment
• set the partial_parse flag

Does that sound right?

[dwreeves]

1. run dbt parse locally / in CICD (pre-deploying), which generates the necessary artifacts

Yes or dbt compile. Both of these commands will produce a manifest.json and partial_parse.msgpack. (A note: This code addition only uses partial_parse.msgpack; however a user making use of pre-compilation for performance reasons would also probably want to use manifest.json as a load method similarly for performance reasons. So two birds with one stone.)

2. push the entire project to the Airflow deployment
3. set the partial_parse flag

By default, partial_parse is True. They wouldn't need to explicitly set it unless they wanted to.

I made this API decision (i.e. the decision to have default as True) in part because the user already receives the following warning from dbt when they are not using partial parsing:

Unable to do partial parsing because saved manifest not found. Starting full parse.

Meaning that a user who intended on using partial parsing, and one who did not, is already being warned by dbt that they are not; and by default dbt is looking for it. So I think having the default behavior of the operator match the default behavior of dbt is sensible. (Also, the performance implications of checking for the existence of the file, and failing to find it, are negligible, so very low harm in doing so.)

That's a lot of text for a very small API decision, but I think it's really important to get stuff like this right for popular packages. And I do think this is the right call.

---

I agree the performant mode is nice in theory but may be difficult to maintain. Maybe we should create a docs page specifically on performance with our recommendations? Inclusive of this flag, using the manifest parsing method, install LibYAML, etc. That way, for users who care about performance, there's a single doc for all optimization recommendations.

I do like this approach. I agree that this would be better than a "performant" mode.

[dwreeves]

I made updates as per everyone's suggestions. Thanks.

The docs will have to make it clear that it's only used for rendering when dbt_ls is used and execution for local or virtualenv.

A note on the documentation. I am not documenting the behavior for LoadMode.DBT_LS. I think, in most cases, it will be rare for a user to both be using DBT_LS and have a partial_parse.msgpack. This would not be completely unheard of and can be a reasonable deliberate decision if the dbt project mutates a lot, but the partial_parse.msgpack is not being updated frequently, so you either set the load mode deliberately or just forego the manifest.json but keep the partial parse file. But all told this is a pretty niche situation-- I think a lot of users on DBT_LS mode are likely foregoing the use of dbt artifacts altogether. And I wasn't sure where it would fit in without feeling like it was adding clutter and potential confusion to the docs for most users. Let me know if you want to document this regardless, but that was my rationale for not mentioning it.

[jbandoro]

Thanks, this is great @dwreeves!

Could you also add the partial_parse arg to the project config docs here: https://github.com/astronomer/astronomer-cosmos/blob/main/docs/configuration/project-config.rst

After that and the minor comment below, it looks good to me to merge.

[dwreeves]

@jbandoro Done