Add node config to TaskInstance Context by linchun3 · Pull Request #1044 · astronomer/astronomer-cosmos

linchun3 · 2024-06-14T14:47:50Z

Description

Add the node's attributes (config, tags, etc, ...) into a TaskInstance context for retrieval by callback functions in Airflow through the use of pre_execute to store these attributes into a task's context.

As this PR seems to be closed, and I have a use case for this feature, I attempt to recreate the needed feature.

We leverage the context_merge utility function from Airflow to merge the extra context into the Context object of a TaskInstance.

Related Issue(s)

closes #698

Breaking Change?

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 · 2024-06-14T14:48:16Z

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

Name	Link
🔨 Latest commit	`62cf7f3`
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/66794283cd8f4400079184be

codecov · 2024-06-17T07:09:22Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.82%. Comparing base (3e38af9) to head (251df0c).
Report is 2 commits behind head on main.

❗ Current head 251df0c differs from pull request most recent head 62cf7f3

Please upload reports for the commit 62cf7f3 to get more accurate results.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1044      +/-   ##
==========================================
+ Coverage   95.81%   95.82%   +0.01%     
==========================================
  Files          62       62              
  Lines        3010     3020      +10     
==========================================
+ Hits         2884     2894      +10     
  Misses        126      126

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

tatiana

Hi @linchun3 thanks a lot for the context and implementing this feature and tests! It makes lots of sense.

I see that it is currently in draft mode, please, move out of draft once you want it to be reviewed :)

tatiana

Hi @linchun3 I left a minor question - we're close to merging this and I think we will be able to ship it in 1.5.

Add a node's config into TaskInstance context so that we can retrieve it and use it for airflow callbacks

This would allow the user to interact with all attributes of a DbtNode

tatiana

Thanks a lot for implementing this and addressing the feedback, @linchun3 ! It looks great!

@tatiana

New Features * Speed up ``LoadMode.DBT_LS`` by caching dbt ls output in Airflow Variable by @tatiana in #1014 * Support to cache profiles created via ``ProfileMapping`` by @pankajastro in #1046 * Support for running dbt tasks in AWS EKS in #944 by @VolkerSchiewe * Add Clickhouse profile mapping by @roadan and @pankajastro in #353 and #1016 * Add node config to TaskInstance Context by @linchun3 in #1044 Bug fixes * Support partial parsing when cache is disabled by @tatiana in #1070 * Fix disk permission error in restricted env by @pankajastro in #1051 * Add CSP header to iframe contents by @dwreeves in #1055 * Stop attaching log adaptors to root logger to reduce logging costs by @glebkrapivin in #1047 Enhancements * Support ``static_index.html`` docs by @dwreeves in #999 * Support deep linking dbt docs via Airflow UI by @dwreeves in #1038 * Add ability to specify host/port for Snowflake connection by @whummer in #1063 Docs * Fix rendering for env ``enable_cache_dbt_ls`` by @pankajastro in #1069 Others * Update documentation for DbtDocs generator by @arjunanan6 in #1043 * Use uv in CI by @dwreeves in #1013 * Cache hatch folder in the CI by @tatiana in #1056 * Change example DAGs to use ``example_conn`` as opposed to ``airflow_db`` by @tatiana in #1054 * Mark plugin integration tests as integration by @tatiana in #1057 * Ensure compliance with linting rule D300 by using triple quotes for docstrings by @pankajastro in #1049 * Pre-commit hook updates in #1039, #1050, #1064 * Remove duplicates in changelog by @jedcunningham in #1068

Add the node's attributes (config, tags, etc, ...) into a TaskInstance context for retrieval by callback functions in Airflow through the use of `pre_execute` to store these attributes into a task's context. As [this PR](https://github.com/astronomer/astronomer-cosmos/pull/700/files) seems to be closed, and I have a use case for this feature, I attempt to recreate the needed feature. We leverage the `context_merge` utility function from Airflow to merge the extra context into the `Context` object of a `TaskInstance`. Closes astronomer#698

@tatiana

New Features * Speed up ``LoadMode.DBT_LS`` by caching dbt ls output in Airflow Variable by @tatiana in astronomer#1014 * Support to cache profiles created via ``ProfileMapping`` by @pankajastro in astronomer#1046 * Support for running dbt tasks in AWS EKS in astronomer#944 by @VolkerSchiewe * Add Clickhouse profile mapping by @roadan and @pankajastro in astronomer#353 and astronomer#1016 * Add node config to TaskInstance Context by @linchun3 in astronomer#1044 Bug fixes * Support partial parsing when cache is disabled by @tatiana in astronomer#1070 * Fix disk permission error in restricted env by @pankajastro in astronomer#1051 * Add CSP header to iframe contents by @dwreeves in astronomer#1055 * Stop attaching log adaptors to root logger to reduce logging costs by @glebkrapivin in astronomer#1047 Enhancements * Support ``static_index.html`` docs by @dwreeves in astronomer#999 * Support deep linking dbt docs via Airflow UI by @dwreeves in astronomer#1038 * Add ability to specify host/port for Snowflake connection by @whummer in astronomer#1063 Docs * Fix rendering for env ``enable_cache_dbt_ls`` by @pankajastro in astronomer#1069 Others * Update documentation for DbtDocs generator by @arjunanan6 in astronomer#1043 * Use uv in CI by @dwreeves in astronomer#1013 * Cache hatch folder in the CI by @tatiana in astronomer#1056 * Change example DAGs to use ``example_conn`` as opposed to ``airflow_db`` by @tatiana in astronomer#1054 * Mark plugin integration tests as integration by @tatiana in astronomer#1057 * Ensure compliance with linting rule D300 by using triple quotes for docstrings by @pankajastro in astronomer#1049 * Pre-commit hook updates in astronomer#1039, astronomer#1050, astronomer#1064 * Remove duplicates in changelog by @jedcunningham in astronomer#1068

@tatiana

New Features * Speed up ``LoadMode.DBT_LS`` by caching dbt ls output in Airflow Variable by @tatiana in #1014 * Support to cache profiles created via ``ProfileMapping`` by @pankajastro in #1046 * Support for running dbt tasks in AWS EKS in #944 by @VolkerSchiewe * Add Clickhouse profile mapping by @roadan and @pankajastro in #353 and #1016 * Add node config to TaskInstance Context by @linchun3 in #1044 Bug fixes * Support partial parsing when cache is disabled by @tatiana in #1070 * Fix disk permission error in restricted env by @pankajastro in #1051 * Add CSP header to iframe contents by @dwreeves in #1055 * Stop attaching log adaptors to root logger to reduce logging costs by @glebkrapivin in #1047 Enhancements * Support ``static_index.html`` docs by @dwreeves in #999 * Support deep linking dbt docs via Airflow UI by @dwreeves in #1038 * Add ability to specify host/port for Snowflake connection by @whummer in #1063 Docs * Fix rendering for env ``enable_cache_dbt_ls`` by @pankajastro in #1069 Others * Update documentation for DbtDocs generator by @arjunanan6 in #1043 * Use uv in CI by @dwreeves in #1013 * Cache hatch folder in the CI by @tatiana in #1056 * Change example DAGs to use ``example_conn`` as opposed to ``airflow_db`` by @tatiana in #1054 * Mark plugin integration tests as integration by @tatiana in #1057 * Ensure compliance with linting rule D300 by using triple quotes for docstrings by @pankajastro in #1049 * Pre-commit hook updates in #1039, #1050, #1064 * Remove duplicates in changelog by @jedcunningham in #1068 (cherry picked from commit 18d2c90)

linchun3 had a problem deploying to external June 14, 2024 14:47 — with GitHub Actions Error

linchun3 force-pushed the add-node-config-to-instance-task-context branch from d913290 to 5525478 Compare June 14, 2024 15:06

linchun3 had a problem deploying to external June 14, 2024 15:06 — with GitHub Actions Error

linchun3 force-pushed the add-node-config-to-instance-task-context branch from 5525478 to e877a6d Compare June 15, 2024 03:14

linchun3 had a problem deploying to external June 15, 2024 03:14 — with GitHub Actions Error

linchun3 force-pushed the add-node-config-to-instance-task-context branch from e877a6d to d4b50e6 Compare June 15, 2024 04:23

linchun3 temporarily deployed to external June 15, 2024 04:24 — with GitHub Actions Inactive

linchun3 mentioned this pull request Jun 15, 2024

Draft: Add config to TaskInstance context #700

Closed

2 tasks

tatiana reviewed Jun 17, 2024

View reviewed changes

linchun3 force-pushed the add-node-config-to-instance-task-context branch from d4b50e6 to 72f0ca9 Compare June 18, 2024 14:32

linchun3 had a problem deploying to external June 18, 2024 14:32 — with GitHub Actions Error

linchun3 force-pushed the add-node-config-to-instance-task-context branch from 72f0ca9 to 54cffca Compare June 18, 2024 14:33

linchun3 had a problem deploying to external June 18, 2024 14:34 — with GitHub Actions Error

linchun3 marked this pull request as ready for review June 18, 2024 14:34

dosubot Bot added the size:L This PR changes 100-499 lines, ignoring generated files. label Jun 18, 2024

linchun3 requested a review from tatiana June 18, 2024 14:34

linchun3 force-pushed the add-node-config-to-instance-task-context branch from 54cffca to d1c6896 Compare June 18, 2024 14:38

linchun3 temporarily deployed to external June 18, 2024 14:38 — with GitHub Actions Inactive

tatiana reviewed Jun 20, 2024

View reviewed changes

Comment thread cosmos/airflow/graph.py Outdated

tatiana reviewed Jun 20, 2024

View reviewed changes

tatiana added this to the Cosmos 1.5.0 milestone Jun 20, 2024

linchun3 had a problem deploying to external June 20, 2024 17:56 — with GitHub Actions Error

linchun3 force-pushed the add-node-config-to-instance-task-context branch from 1fa1e91 to 2a73a3d Compare June 20, 2024 17:58

linchun3 had a problem deploying to external June 20, 2024 17:59 — with GitHub Actions Error

linchun3 requested a review from tatiana June 20, 2024 18:00

linchun3 added 2 commits June 23, 2024 08:17

Add node config to TaskInstance Context

d93aff1

Add a node's config into TaskInstance context so that we can retrieve it and use it for airflow callbacks

Update feature to insert attributes of DbtNode into extra_context

251df0c

This would allow the user to interact with all attributes of a DbtNode

linchun3 force-pushed the add-node-config-to-instance-task-context branch from 2a73a3d to 251df0c Compare June 23, 2024 00:17

linchun3 temporarily deployed to external June 23, 2024 00:18 — with GitHub Actions Inactive

Merge branch 'main' into add-node-config-to-instance-task-context

62cf7f3

tatiana had a problem deploying to external June 24, 2024 09:55 — with GitHub Actions Failure

tatiana approved these changes Jun 24, 2024

View reviewed changes

dosubot Bot added the lgtm This PR has been approved by a maintainer label Jun 24, 2024

tatiana merged commit 5f5616c into astronomer:main Jun 24, 2024

pankajkoti mentioned this pull request Jun 27, 2024

Release 1.5.0 #1071

Merged

tatiana mentioned this pull request Jun 27, 2024

New Features #984

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add node config to TaskInstance Context#1044

Add node config to TaskInstance Context#1044
tatiana merged 3 commits into
astronomer:mainfrom
linchun3:add-node-config-to-instance-task-context

linchun3 commented Jun 14, 2024 •

edited

Loading

Uh oh!

netlify Bot commented Jun 14, 2024 •

edited

Loading

Uh oh!

codecov Bot commented Jun 17, 2024 •

edited

Loading

Uh oh!

tatiana left a comment

Uh oh!

Uh oh!

tatiana left a comment

Uh oh!

tatiana left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

linchun3 commented Jun 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue(s)

Breaking Change?

Checklist

Uh oh!

netlify Bot commented Jun 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Uh oh!

codecov Bot commented Jun 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

tatiana left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tatiana left a comment

Choose a reason for hiding this comment

Uh oh!

tatiana left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

linchun3 commented Jun 14, 2024 •

edited

Loading

netlify Bot commented Jun 14, 2024 •

edited

Loading

codecov Bot commented Jun 17, 2024 •

edited

Loading