diff --git a/libs/client/cornflow_client/airflow/dag_utilities.py b/libs/client/cornflow_client/airflow/dag_utilities.py index 6e5a4ae5..741b148e 100644 --- a/libs/client/cornflow_client/airflow/dag_utilities.py +++ b/libs/client/cornflow_client/airflow/dag_utilities.py @@ -78,7 +78,7 @@ def connect_to_cornflow(secrets): """ # This secret comes from airflow configuration print("Getting connection information from ENV VAR=CF_URI") - uri = secrets.get_conn_uri("CF_URI") + uri = secrets.get_conn_value("CF_URI") conn = urlparse(uri) scheme = conn.scheme if scheme == "cornflow": @@ -326,7 +326,7 @@ def callback_email(context): notification_email = EnvironmentVariablesBackend().get_variable( "NOTIFICATION_EMAIL" ) - environment_name = EnvironmentVariablesBackend().get_variable("ENVIRONMENT_NAME") + environment_name = os.getenv("AIRFLOW__WEBSERVER__INSTANCE_NAME", "CornflowEnv") title = f"Airflow. {environment_name} ({environment}). DAG/task error: {context['dag'].dag_id}/{context['ti'].task_id} Failed" body = f""" diff --git a/libs/client/cornflow_client/tests/unit/test_dag_utilities.py b/libs/client/cornflow_client/tests/unit/test_dag_utilities.py index fdaa33ab..6469dd85 100644 --- a/libs/client/cornflow_client/tests/unit/test_dag_utilities.py +++ b/libs/client/cornflow_client/tests/unit/test_dag_utilities.py @@ -31,8 +31,8 @@ def test_env_connection_vars(self, CornFlow): ] client_instance = CornFlow.return_value client_instance.login.return_value = "" - for (conn_str, user_info, url) in conn_uris: - secrets.get_conn_uri.return_value = conn_str + for conn_str, user_info, url in conn_uris: + secrets.get_conn_value.return_value = conn_str du.connect_to_cornflow(secrets) client_instance.login.assert_called_with( username=user_info[0], pwd=user_info[1]