From e6f79bff64cf40ce0418e72d7ceba48c2f9f6461 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Fri, 27 Jun 2025 19:06:25 +0530 Subject: [PATCH] Replace usage of 'set_extra' with 'extra' for athena sql hook --- .../src/airflow/providers/amazon/aws/hooks/athena_sql.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py index e8d8680c29c11..d66f788226aa8 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena_sql.py @@ -111,7 +111,14 @@ def conn_config(self) -> AwsConnectionWrapper: connection.login = athena_conn.login connection.password = athena_conn.password connection.schema = athena_conn.schema - connection.set_extra(json.dumps({**athena_conn.extra_dejson, **connection.extra_dejson})) + merged_extra = {**athena_conn.extra_dejson, **connection.extra_dejson} + try: + extra_json = json.dumps(merged_extra) + connection.extra = extra_json + except (TypeError, ValueError): + raise ValueError( + f"Encountered non-JSON in `extra` field for connection {self.aws_conn_id!r}." + ) except AirflowNotFoundException: connection = athena_conn connection.conn_type = "aws"