flowrs-tui 0.1.15

Flowrs is a Terminal User Interface (TUI) for Apache Airflow
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

use anyhow::Result;
use log::error;
use reqwest::{Method, Response};

use super::AirFlowClient;
use crate::airflow::model::dag::DagList;

impl AirFlowClient {
    pub async fn list_dags(&self) -> Result<DagList> {
        let r = self.base_api(Method::GET, "dags")?.build()?;
        let response = self.client.execute(r).await?;

        let daglist = response.json::<DagList>().await;
        match daglist {
            Ok(daglist) => Ok(daglist),
            Err(e) => {
                error!("Error parsing response: {}", e);
                Ok(DagList::default())
            }
        }
    }

    // pub async fn get_dag_details(&self, dag_id: &str) -> Result<()> {
    //     let r = self
    //         .base_api(Method::GET, &format!("dags/{dag_id}/details"))?
    //         .build()?;
    //     let response = self.client.execute(r).await?;
    //     let _ = response.text().await?;
    //     Ok(())
    // }

    pub async fn toggle_dag(&self, dag_id: &str, is_paused: bool) -> Result<()> {
        let _: Response = self
            .base_api(Method::PATCH, &format!("dags/{dag_id}"))?
            .query(&[("update_mask", "is_paused")])
            .json(&serde_json::json!({"is_paused": !is_paused}))
            .send()
            .await?;
        Ok(())
    }

    pub async fn get_dag_code(&self, file_token: &String) -> Result<String> {
        let r = self
            .base_api(Method::GET, &format!("dagSources/{file_token}"))?
            .build()?;
        let response = self.client.execute(r).await?;
        let code = response.text().await?;
        Ok(code)
    }
}

#[cfg(test)]
mod tests {

    use super::AirFlowClient;
    use super::DagList;
    use crate::airflow::config::FlowrsConfig;

    const TEST_CONFIG: &str = r#"[[servers]]
        name = "test"
        endpoint = "http://localhost:8080"

        [servers.auth.BasicAuth]
        username = "airflow"
        password = "airflow"
        "#;

    #[tokio::test]
    async fn test_list_dags() {
        let config: FlowrsConfig = toml::from_str(str::trim(TEST_CONFIG)).unwrap();
        let client = AirFlowClient::new(config.servers.unwrap()[0].clone()).unwrap();
        let daglist: DagList = client.list_dags().await.unwrap();
        assert_eq!(daglist.dags[0].owners, vec!["airflow"]);
    }

    #[tokio::test]
    async fn test_get_dag_code() {
        let config: FlowrsConfig = toml::from_str(str::trim(TEST_CONFIG)).unwrap();
        let client = AirFlowClient::new(config.servers.unwrap()[0].clone()).unwrap();

        let dag = client.list_dags().await.unwrap().dags[0].clone();
        let code = client.get_dag_code(&dag.file_token).await.unwrap();
        assert!(code.contains("with DAG"));
    }
}