Skip to content
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.

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

Fix aws assume role session creation when deferrable #40051

Conversation

gopidesupavan
Copy link
Contributor

Updated _create_session_with_assume_role to return aiobotocore.session.AioSession if session requested with deferrable=True

Tested with example:

aws_assume_role_connection = {
"region_name": "us-east-1",
"role_arn": "arn:aws:iam::{ACCOUNT_ID}:role/TestRoleAirflow",
"assume_role_method": "assume_role"
}

with DAG(
    dag_id=DAG_ID,
    start_date=datetime(2021, 1, 1),
    schedule=None,
    catchup=False,
) as dag:
    bucket = "s3://environ-issues-fix-repo/"

    base = ObjectStoragePath(bucket, conn_id="aws_assume_role_connection")

    key = f"test/sample.json"


    @task
    def read_file(path: ObjectStoragePath) -> None:
        with path.open() as f:
            print(json.loads(f.read()))


    read_file(base / key)
image

Closes: #39970


^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

@gopidesupavan
Copy link
Contributor Author

cc: @Taragolis

@eladkal eladkal requested review from ferruzzi and vincbeck June 6, 2024 03:12
@ferruzzi ferruzzi merged commit 42a2b1a into apache:main Jun 6, 2024
49 checks passed
@gopidesupavan gopidesupavan deleted the fix-aws-assume-role-session-with-deferrable branch July 5, 2024 12:29
romsharon98 pushed a commit to romsharon98/airflow that referenced this pull request Jul 26, 2024
* fixed aws connection issue for  assume_role when session created with defarrable=True
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
area:providers provider:amazon-aws AWS/Amazon - related issues
Projects
None yet
Development

Successfully merging this pull request may close these issues.

S3 ObjectStorage can't use AWS connection with role
3 participants