Skip to content

Setting Airflow context Environment variables for operators#47644

Merged
amoghrajesh merged 20 commits into
apache:mainfrom
astronomer:python_op_export_env_vars
Mar 21, 2025
Merged

Setting Airflow context Environment variables for operators#47644
amoghrajesh merged 20 commits into
apache:mainfrom
astronomer:python_op_export_env_vars

Conversation

@prabhusneha

@prabhusneha prabhusneha commented Mar 11, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Closes: #46418

This change passes down the environment variables through context.

Testing:

DAG used:

from datetime import timedelta
import requests
import time
import os

from pendulum import today

from airflow import DAG
from airflow.operators.python import PythonOperator

def log_checker(ti_id: str, expected: str, notexpected: str, try_number: int = 1, **context: dict):
    time.sleep(30)
    dag_instance = context['dag']
    dagrun = context["dag_run"]
    dag_id = dagrun.dag_id
    run_id = dagrun.run_id
    task_id = dag_instance.get_task(ti_id).task_id
    print("Printing env vars: owner")
    print(os.environ.get('AIRFLOW_CTX_DAG_OWNER'))
    
    
def easy_return():
    pass


default_args = {"owner": "airflow_3_0_0"}

with DAG(
    dag_id="test_setting_context",
    start_date=today('UTC').add(days=-1),
    schedule=None,
    default_args=default_args,
    tags=["dagparams"],
) as dag:

    py0 = PythonOperator(task_id="dummy1", python_callable=easy_return)

    py1 = PythonOperator(
        task_id="check_for_owner_in_def_args",
        python_callable=log_checker,
        op_args=[
            "dummy1",
            "AIRFLOW_CTX_DAG_OWNER='airflow'",
            "AIRFLOW_CTX_DAG_OWNER='astro'",
        ],
    )
py0 >> py1

We will now need to print/log the environment variables in the dag in Airflow 3. This PR only sets the airflow context environment variables.

image

Running the same DAG in Airflow 2.10.5, give the log as seen below:

image

@prabhusneha prabhusneha mentioned this pull request Mar 11, 2025
Closed
2 tasks
@phanikumv phanikumv requested review from amoghrajesh and kaxil March 12, 2025 04:55
@vatsrahul1001

vatsrahul1001 commented Mar 12, 2025

Copy link
Copy Markdown
Contributor

Can we add a test for this?

@amoghrajesh

amoghrajesh commented Mar 12, 2025

Copy link
Copy Markdown
Contributor

@prabhusneha can we please add some unit tests and also test with the dag that was originally provided in the bug report?

ashb
ashb reviewed Mar 12, 2025
View reviewed changes
Comment thread task-sdk/src/airflow/sdk/execution_time/task_runner.py Outdated
Comment thread task-sdk/src/airflow/sdk/execution_time/task_runner.py Outdated
@prabhusneha

prabhusneha commented Mar 13, 2025

Copy link
Copy Markdown
Contributor Author

Added unit tests

@prabhusneha

prabhusneha commented Mar 13, 2025

Copy link
Copy Markdown
Contributor Author

Fixing the failing CI checks

Lee-W
Lee-W reviewed Mar 14, 2025
View reviewed changes
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py Outdated
amoghrajesh
amoghrajesh reviewed Mar 14, 2025
View reviewed changes
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py Outdated
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py Outdated
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
Comment thread task-sdk/src/airflow/sdk/execution_time/task_runner.py Outdated
Comment thread task-sdk/tests/task_sdk/execution_time/test_task_runner.py Outdated
@prabhusneha prabhusneha force-pushed the python_op_export_env_vars branch from 1d32ff2 to c30c4d8 Compare March 18, 2025 11:48
amoghrajesh
amoghrajesh reviewed Mar 20, 2025
View reviewed changes
Comment thread airflow/version_compat.py
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/version_compat.py
Comment thread providers/apache/hive/tests/unit/apache/hive/transfers/test_hive_to_mysql.py
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py Outdated
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py Outdated
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py Outdated
Comment thread task-sdk/tests/task_sdk/execution_time/test_task_runner.py Outdated
@amoghrajesh

amoghrajesh commented Mar 20, 2025

Copy link
Copy Markdown
Contributor

There's a static check failure too:

mypy providers/http/src providers/http/tests providers/sqlite/src providers/sqlite/tests providers/docker/src providers/docker/tests providers/jenkins/src providers/jenkins/tests providers/celery/src providers/celery/tests providers/postgres/src providers/postgres/tests providers/amazon/src providers/amazon/tests providers/opsgenie/src providers/opsgenie/tests providers/redis/src providers/redis/tests providers/smtp/src providers/smtp/tests providers/sftp/src providers/sftp/tests providers/samba/src providers/samba/tests providers/snowflake/src providers/snowflake/tests providers/alibaba/src providers/alibaba/tests providers/tableau/src providers/tableau/tests providers/slack/src providers/slack/tests providers/salesforce/src providers/salesforce/tests providers/yandex/src providers/yandex/tests providers/presto/src providers/presto/tests providers/vertica/src providers/vertica/tests providers/cncf/kubernetes/src providers/cncf/kubernetes/tests providers/opensearch/src providers/opensearch/tests providers/odbc/src providers/odbc/tests providers/openfaas/src providers/openfaas/tests providers/trino/src providers/trino/tests providers/arangodb/src providers/arangodb/tests providers/hashicorp/src providers/hashicorp/tests providers/mongo/src providers/mongo/tests providers/sendgrid/src providers/sendgrid/tests providers/papermill/src providers/papermill/tests providers/facebook/src providers/facebook/tests providers/standard/src providers/standard/tests providers/apprise/src providers/apprise/tests providers/pgvector/src providers/pgvector/tests providers/dingding/src providers/dingding/tests providers/mysql/src providers/mysql/tests providers/zendesk/src providers/zendesk/tests providers/google/src providers/google/tests providers/ssh/src providers/ssh/tests providers/exasol/src providers/exasol/tests providers/grpc/src providers/grpc/tests providers/influxdb/src providers/influxdb/tests providers/dbt/cloud/src providers/dbt/cloud/tests providers/segment/src providers/segment/tests providers/cohere/src providers/cohere/tests providers/edge/src providers/edge/tests providers/cloudant/src providers/cloudant/tests providers/fab/src providers/fab/tests providers/pagerduty/src providers/pagerduty/tests providers/jdbc/src providers/jdbc/tests providers/telegram/src providers/telegram/tests providers/weaviate/src providers/weaviate/tests providers/ydb/src providers/ydb/tests providers/common/io/src providers/common/io/tests providers/common/compat/src providers/common/compat/tests providers/common/messaging/src providers/common/messaging/tests providers/common/sql/src providers/common/sql/tests providers/elasticsearch/src providers/elasticsearch/tests providers/asana/src providers/asana/tests providers/airbyte/src providers/airbyte/tests providers/ftp/src providers/ftp/tests providers/openlineage/src providers/openlineage/tests providers/github/src providers/github/tests providers/imap/src providers/imap/tests providers/neo4j/src providers/neo4j/tests providers/discord/src providers/discord/tests providers/atlassian/jira/src providers/atlassian/jira/tests providers/apache/druid/src providers/apache/druid/tests providers/apache/drill/src providers/apache/drill/tests providers/apache/impala/src providers/apache/impala/tests providers/apache/pinot/src providers/apache/pinot/tests providers/apache/kylin/src providers/apache/kylin/tests providers/apache/spark/src providers/apache/spark/tests providers/apache/cassandra/src providers/apache/cassandra/tests providers/apache/iceberg/src providers/apache/iceberg/tests providers/apache/livy/src providers/apache/livy/tests providers/apache/hdfs/src providers/apache/hdfs/tests providers/apache/kafka/src providers/apache/kafka/tests providers/apache/hive/src providers/apache/hive/tests providers/apache/pig/src providers/apache/pig/tests providers/apache/beam/src providers/apache/beam/tests providers/apache/flink/src providers/apache/flink/tests providers/databricks/src providers/databricks/tests providers/datadog/src providers/datadog/tests providers/teradata/src providers/teradata/tests providers/openai/src providers/openai/tests providers/qdrant/src providers/qdrant/tests providers/oracle/src providers/oracle/tests providers/pinecone/src providers/pinecone/tests providers/microsoft/mssql/src providers/microsoft/mssql/tests providers/microsoft/winrm/src providers/microsoft/winrm/tests providers/microsoft/azure/src providers/microsoft/azure/tests providers/microsoft/psrp/src providers/microsoft/psrp/tests providers/singularity/src providers/singularity/tests --namespace-packages
providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py:46: error:
Module "airflow.utils.operator_helpers" has no attribute
"AIRFLOW_VAR_NAME_FORMAT_MAPPING"  [attr-defined]
        from airflow.utils.operator_helpers import (
        ^
providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py:46: error:
Name "AIRFLOW_VAR_NAME_FORMAT_MAPPING" already defined (possibly by an import) 
[no-redef]
        from airflow.utils.operator_helpers import (

amoghrajesh
amoghrajesh reviewed Mar 21, 2025
View reviewed changes
Comment thread airflow/models/taskinstance.py Outdated
amoghrajesh
amoghrajesh reviewed Mar 21, 2025
View reviewed changes

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some nits about compat. Should fix the CI

Comment thread generated/provider_dependencies.json
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/hooks/hive.py
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/operators/hive.py
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_mysql.py
Comment thread providers/apache/hive/src/airflow/providers/apache/hive/transfers/hive_to_samba.py
Comment thread providers/apache/hive/tests/unit/apache/hive/hooks/test_hive.py Outdated
Comment thread providers/presto/src/airflow/providers/presto/hooks/presto.py
Comment thread providers/standard/src/airflow/providers/standard/operators/bash.py
Comment thread providers/trino/src/airflow/providers/trino/hooks/trino.py
@prabhusneha

prabhusneha commented Mar 21, 2025

Copy link
Copy Markdown
Contributor Author

@amoghrajesh Please take a look. CI checks are passing.

@amoghrajesh amoghrajesh self-requested a review March 21, 2025 09:20
amoghrajesh
amoghrajesh approved these changes Mar 21, 2025
View reviewed changes

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! LGTM +1

@amoghrajesh

amoghrajesh commented Mar 21, 2025

Copy link
Copy Markdown
Contributor

This will require a rebase, let me take a stab at that

@amoghrajesh amoghrajesh merged commit 545bf45 into apache:main Mar 21, 2025
@amoghrajesh amoghrajesh deleted the python_op_export_env_vars branch March 21, 2025 15:23
agupta01 pushed a commit to agupta01/airflow that referenced this pull request Mar 21, 2025
@amoghrajesh @agupta01
Setting Airflow context Environment variables for operators (apache#4…
3c9c002 
…7644)

Co-authored-by: Sneha Prabhu <snehaprabhu@Snehas-MacBook-Pro.local>
Co-authored-by: Amogh Desai <amoghrajesh1999@gmail.com>
shubham-pyc pushed a commit to shubham-pyc/airflow that referenced this pull request Mar 22, 2025
@amoghrajesh @shubham-pyc
Setting Airflow context Environment variables for operators (apache#4…
c3f6d38 
…7644)

Co-authored-by: Sneha Prabhu <snehaprabhu@Snehas-MacBook-Pro.local>
Co-authored-by: Amogh Desai <amoghrajesh1999@gmail.com>
@atul-astronomer

atul-astronomer commented Mar 24, 2025

Copy link
Copy Markdown
Contributor

As per the description, now env variables are not printed in the logs explicitly as it was in AF2 but if user wants, they can print it using 'print(os.environ.get('AIRFLOW_CTX_DAG_OWNER'))'

@eladkal eladkal mentioned this pull request Mar 26, 2025
Closed
nailo2c pushed a commit to nailo2c/airflow that referenced this pull request Apr 4, 2025
@amoghrajesh @nailo2c
Setting Airflow context Environment variables for operators (apache#4…
b7d9254 
…7644)

Co-authored-by: Sneha Prabhu <snehaprabhu@Snehas-MacBook-Pro.local>
Co-authored-by: Amogh Desai <amoghrajesh1999@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ashb ashb ashb left review comments
@Lee-W Lee-W Lee-W left review comments
@amoghrajesh amoghrajesh amoghrajesh approved these changes
@kaxil kaxil Awaiting requested review from kaxil kaxil is a code owner
@jscheffl jscheffl Awaiting requested review from jscheffl
@XD-DENG XD-DENG Awaiting requested review from XD-DENG XD-DENG is a code owner

Assignees

No one assigned

Labels

area:task-sdk

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Airflow context Env variables are not set when using PythonOperator

7 participants

@prabhusneha @vatsrahul1001 @amoghrajesh @atul-astronomer @ashb @Lee-W @kaxil