diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py index 34d8673d18a96..f90de71855f9d 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/hooks/kubernetes.py @@ -140,7 +140,8 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: def __init__( self, - conn_id: str | None = default_conn_name, + conn_id: str | None = None, + kubernetes_conn_id: str | None = default_conn_name, client_configuration: client.Configuration | None = None, cluster_context: str | None = None, config_file: str | None = None, @@ -149,7 +150,7 @@ def __init__( disable_tcp_keepalive: bool | None = None, ) -> None: super().__init__() - self.conn_id = conn_id + self.conn_id = conn_id or kubernetes_conn_id self.client_configuration = client_configuration self.cluster_context = cluster_context self.config_file = config_file @@ -706,6 +707,14 @@ def get_yaml_content_from_file(kueue_yaml_url) -> list[dict]: return list(yaml.safe_load_all(response.text)) + def test_connection(self): + try: + conn = self.get_conn() + version: client.VersionInfo = client.VersionApi(conn).get_code() + return True, f"Connection successful. Version Info: {version.to_dict()}" + except Exception as e: + return False, str(e) + def _get_bool(val) -> bool | None: """Convert val to bool if can be done with certainty; if we cannot infer intention we return None."""