ashb commented on code in PR #53831:
URL: https://github.com/apache/airflow/pull/53831#discussion_r2249148592
##########
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:
+ """
+ Get connection (async), given connection id.
+
+ :param conn_id: connection id
+ :return: connection
+ """
+ if not cls._is_task_sdk_context():
+ raise AirflowException("This method currently only supports
getting connections through Task SDK")
Review Comment:
```suggestion
raise RuntimeException("This method currently only supports
getting connections through Task SDK")
```
AirflowException is over used and has no special meaning, so I think
RuntimeException is better suited here.
--
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]
