Skip to content

AIP-72: Port Variable.set From TaskSDK to Models #48177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
jason810496 wants to merge 6 commits into from
Closed

AIP-72: Port Variable.set From TaskSDK to Models #48177

jason810496 wants to merge 6 commits into from

Conversation

@jason810496
Copy link
Member

@jason810496 jason810496 commented Mar 24, 2025 •
edited
Loading

closes: #47920

Why

Setting a variable in Dag is failing due to 'Direct database access via the ORM is not allowed in Airflow 3.0' ( details in issue context )

from datetime import datetime

from airflow.models import Variable
from airflow import DAG

my_var = Variable.set("param_variable", 10)

What

Port Variable.set from TaskSDK to airflow.models.Variable.set

@jason810496 jason810496 force-pushed the feature/task-sdk/port-variable-set-method branch 3 times, most recently from 49268f9 to 51a3a7e Compare March 25, 2025 07:44
@jason810496 jason810496 marked this pull request as ready for review March 25, 2025 07:49
@jason810496 jason810496 force-pushed the feature/task-sdk/port-variable-set-method branch from 51a3a7e to a076fa4 Compare March 25, 2025 07:49
jason810496
jason810496 commented Mar 25, 2025
View reviewed changes
task-sdk/src/airflow/sdk/execution_time/context.py
Comment on lines +196 to +229
def _set_variable(key: str, value: Any, description: str | None = None, serialize_json: bool = False) -> None:
from airflow.sdk.execution_time.comms import ErrorResponse, PutVariable
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS

if serialize_json:
import json

value = json.dumps(value, indent=2)
else:
value = str(value)

SUPERVISOR_COMMS.send_request(log=log, msg=PutVariable(key=key, value=value, description=description))
msg = SUPERVISOR_COMMS.get_message()
if isinstance(msg, ErrorResponse):
raise AirflowRuntimeError(msg)
return


Copy link
Member Author

@jason810496 jason810496 Mar 25, 2025

Choose a reason for hiding this comment

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

I think I need to port the Secret Backend check_for_write_conflict logic from models ( https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/models/variable.py#L198-L199 ) to TaskSDK as well?

Just like _get_variable, right?
https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/execution_time/context.py#L154-L171

cc @amoghrajesh

Copy link
Contributor

@amoghrajesh amoghrajesh Apr 9, 2025

Choose a reason for hiding this comment

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

Thats right we need to do that as well

@Lee-W Lee-W self-requested a review March 25, 2025 11:34
@jason810496 jason810496 force-pushed the feature/task-sdk/port-variable-set-method branch from a076fa4 to 9f83b4c Compare March 30, 2025 13:49
@jason810496
Copy link
Member Author

jason810496 commented Mar 30, 2025

Just rebased to latest main to resolve conflict.

@jason810496 jason810496 mentioned this pull request Mar 31, 2025
Merged
amoghrajesh
amoghrajesh reviewed Apr 9, 2025
View reviewed changes
Copy link
Contributor

@amoghrajesh amoghrajesh left a comment

Choose a reason for hiding this comment

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

@jason810496 this only handles the top level variables. We will also have to port over / allow this from the task runner level too.

Could you please handle that as well?

task-sdk/src/airflow/sdk/execution_time/context.py
Comment on lines +196 to +229
def _set_variable(key: str, value: Any, description: str | None = None, serialize_json: bool = False) -> None:
from airflow.sdk.execution_time.comms import ErrorResponse, PutVariable
from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS

if serialize_json:
import json

value = json.dumps(value, indent=2)
else:
value = str(value)

SUPERVISOR_COMMS.send_request(log=log, msg=PutVariable(key=key, value=value, description=description))
msg = SUPERVISOR_COMMS.get_message()
if isinstance(msg, ErrorResponse):
raise AirflowRuntimeError(msg)
return


Copy link
Contributor

@amoghrajesh amoghrajesh Apr 9, 2025

Choose a reason for hiding this comment

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

Thats right we need to do that as well

@jason810496 jason810496 force-pushed the feature/task-sdk/port-variable-set-method branch from 9f83b4c to 3ba45df Compare April 9, 2025 07:53
@amoghrajesh
Copy link
Contributor

amoghrajesh commented Apr 9, 2025

Discussed with @jason810496 on this one.

There is more work left to do on this and it would be nice to do this in RC tomorrow eod. I will be working on this parallely and creating a companion PR that handles the missing gaps in this one too.

@jason810496
Copy link
Member Author

jason810496 commented Apr 9, 2025

Closing this one, as #49005 will address the issue.

@jason810496 jason810496 closed this Apr 9, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@amoghrajesh amoghrajesh amoghrajesh left review comments

@ashb ashb Awaiting requested review from ashb ashb is a code owner

@kaxil kaxil Awaiting requested review from kaxil kaxil is a code owner

@jedcunningham jedcunningham Awaiting requested review from jedcunningham jedcunningham is a code owner

@ephraimbuddy ephraimbuddy Awaiting requested review from ephraimbuddy ephraimbuddy is a code owner

@XD-DENG XD-DENG Awaiting requested review from XD-DENG XD-DENG is a code owner

@Lee-W Lee-W Awaiting requested review from Lee-W

Assignees

No one assigned

Labels

area:DAG-processing area:task-execution-interface-aip72 AIP-72: Task Execution Interface (TEI) aka Task SDK area:task-sdk

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Setting a variable in Dag is failing due to 'Direct database access via the ORM is not allowed in Airflow 3.0'

2 participants

@jason810496 @amoghrajesh