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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use log::debug;
use reqwest::{Method, Response};

use crate::airflow::model::dagrun::DagRunList;
use anyhow::Result;

use super::AirFlowClient;

impl AirFlowClient {
    pub async fn list_dagruns(&self, dag_id: &str) -> Result<DagRunList> {
        let response: Response = self
            .base_api(Method::GET, format!("dags/{dag_id}/dagRuns").as_str())?
            .query(&[("order_by", "-execution_date")])
            .send()
            .await?;
        let dagruns: DagRunList = response.json::<DagRunList>().await?;
        Ok(dagruns)
    }

    #[allow(dead_code)]
    pub async fn list_all_dagruns(&self) -> Result<DagRunList> {
        let response: Response = self
            .base_api(Method::POST, "dags/~/dagRuns/list")?
            .json(&serde_json::json!({"page_limit": 200}))
            .send()
            .await?;
        let dagruns: DagRunList = response.json::<DagRunList>().await?;
        Ok(dagruns)
    }

    pub async fn mark_dag_run(&self, dag_id: &str, dag_run_id: &str, status: &str) -> Result<()> {
        let _: Response = self
            .base_api(
                Method::PATCH,
                format!("dags/{dag_id}/dagRuns/{dag_run_id}").as_str(),
            )?
            .json(&serde_json::json!({"state": status}))
            .send()
            .await?;
        Ok(())
    }

    pub async fn clear_dagrun(&self, dag_id: &str, dag_run_id: &str) -> Result<()> {
        let _: Response = self
            .base_api(
                Method::POST,
                format!("dags/{dag_id}/dagRuns/{dag_run_id}/clear").as_str(),
            )?
            .json(&serde_json::json!({"dry_run": false}))
            .send()
            .await?;
        Ok(())
    }

    pub async fn trigger_dag_run(&self, dag_id: &str) -> Result<()> {
        let resp: Response = self
            .base_api(Method::POST, format!("dags/{dag_id}/dagRuns").as_str())?
            .json(&serde_json::json!({}))
            .send()
            .await?;

        debug!("{:?}", resp);
        Ok(())
    }
}

#[cfg(test)]
mod tests {

    use super::AirFlowClient;
    use super::DagRunList;
    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 server = config.servers.unwrap()[0].clone();
        let client = AirFlowClient::new(server).unwrap();
        let first_dag = &client.list_dags().await.unwrap().dags[0];

        println!("{:?}", client.config);
        let dagrun_list: DagRunList = client
            .list_dagruns(first_dag.dag_id.as_str())
            .await
            .unwrap();
        assert!(!dagrun_list.dag_runs.is_empty());
    }

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

        let dagrun_list: DagRunList = client.list_all_dagruns().await.unwrap();
        assert!(!dagrun_list.dag_runs.is_empty());
    }
}