amoghrajesh commented on code in PR #53831:
URL: https://github.com/apache/airflow/pull/53831#discussion_r2250680910
##########
task-sdk/src/airflow/sdk/bases/hook.py:
##########
@@ -56,21 +64,36 @@ def get_connection(cls, conn_id: str) -> Connection:
:param conn_id: connection id
:return: connection
"""
- import sys
-
- # if SUPERVISOR_COMMS is set, we're in task sdk context
- if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"),
"SUPERVISOR_COMMS"):
+ if cls._is_task_sdk_context():
Review Comment:
BTW, this is a temporary layer that going to go off for 3.1 through
https://github.com/apache/airflow/issues/52520, so I do not think we should
introduce a function for it.
##########
task-sdk/src/airflow/sdk/bases/hook.py:
##########
@@ -56,21 +64,36 @@ def get_connection(cls, conn_id: str) -> Connection:
:param conn_id: connection id
:return: connection
"""
- import sys
-
- # if SUPERVISOR_COMMS is set, we're in task sdk context
- if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"),
"SUPERVISOR_COMMS"):
+ if cls._is_task_sdk_context():
from airflow.sdk.definitions.connection import Connection
conn = Connection.get(conn_id)
log.debug("Connection Retrieved '%s' (via task-sdk)", conn.conn_id)
return conn
+
from airflow.models.connection import Connection as ConnectionModel
conn = ConnectionModel.get_connection_from_secrets(conn_id)
log.debug("Connection Retrieved '%s' (via core Airflow)", conn.conn_id)
return conn
+ @classmethod
+ async def async_get_connection(cls, conn_id: str) -> Connection:
Review Comment:
Would be in favour of `aget_connection` too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
