starry/database/
extract.rs

1use {
2    crate::*,
3    anyhow::*,
4    chrono::{
5        DateTime,
6        SecondsFormat,
7        Utc,
8    },
9    cli_log::*,
10    std::{
11        collections::HashMap,
12        io::Write,
13    },
14};
15
16#[derive(Debug)]
17pub struct ExtractLine {
18    pub time: DateTime<Utc>,
19    // this vec is the same len than the extract's names
20    pub counts: Vec<Option<usize>>,
21}
22
23#[derive(Debug)]
24pub struct Extract {
25    // names of either users or repos (with a /)
26    pub names: Vec<String>,
27    pub lines: Vec<ExtractLine>,
28}
29
30#[derive(Debug)]
31pub(crate) struct Col {
32    pub idx: usize,
33    pub name: String,
34}
35
36impl Extract {
37    pub fn write_csv<W: Write>(
38        &self,
39        w: &mut W,
40    ) -> Result<()> {
41        write!(w, "time")?;
42        for name in &self.names {
43            write!(w, ",{}", name)?;
44        }
45        writeln!(w)?;
46        for line in &self.lines {
47            write!(
48                w,
49                "{}",
50                line.time.to_rfc3339_opts(SecondsFormat::Secs, true)
51            )?;
52            for count in &line.counts {
53                if let Some(count) = count {
54                    write!(w, ",{}", count)?;
55                } else {
56                    write!(w, ",")?;
57                }
58            }
59            writeln!(w)?;
60        }
61        w.flush()?;
62        Ok(())
63    }
64    pub fn read(
65        db: &Db,
66        names: Vec<String>,
67    ) -> Result<Self> {
68        // we first compile the user request in several queries (one per user)
69        let mut queries: Vec<UserQuery> = Vec::new();
70        for (idx, name) in names.iter().enumerate() {
71            let mut tokens = name.split('/');
72            let user_id = UserId::new(tokens.next().unwrap()); // SAFETY: first split element is never None
73            let query_idx = queries
74                .iter()
75                .position(|q| q.user_id == user_id)
76                .unwrap_or_else(|| {
77                    let idx = queries.len();
78                    queries.push(UserQuery {
79                        user_id,
80                        sum: None,
81                        repos: Vec::new(),
82                    });
83                    idx
84                });
85            match tokens.next() {
86                Some(repo) => {
87                    queries[query_idx].repos.push(Col {
88                        idx,
89                        name: repo.to_string(),
90                    });
91                }
92                None => {
93                    // if the user specifies twice the same user name, one
94                    // won't be filled. We don't care
95                    queries[query_idx].sum = Some(Col {
96                        idx,
97                        name: name.to_string(),
98                    });
99                }
100            }
101        }
102        debug!("queries: {:#?}", &queries);
103        // we now execute all the queries, storing and merging the results in a map per time
104        let mut results: HashMap<DateTime<Utc>, ExtractLine> = HashMap::new();
105        for query in queries {
106            let repo_names = query.repos.iter().map(|col| col.name.as_str()).collect();
107            let response_lines = db.extract_user_query(&query.user_id, repo_names)?;
108            debug!("response_lines: {:#?}", &response_lines);
109            for response_line in response_lines {
110                let extract_line =
111                    results
112                        .entry(response_line.time)
113                        .or_insert_with(|| ExtractLine {
114                            time: response_line.time,
115                            counts: vec![None; names.len()],
116                        });
117                if let Some(col) = query.sum.as_ref() {
118                    extract_line.counts[col.idx] = Some(response_line.sum);
119                }
120                for (idx, col) in query.repos.iter().enumerate() {
121                    extract_line.counts[col.idx] = response_line.counts[idx];
122                }
123            }
124        }
125        // we sort the lines
126        let mut lines: Vec<ExtractLine> = results.drain().map(|(_, line)| line).collect();
127        debug!("lines: {:#?}", &lines);
128        lines.sort_by_key(|line| line.time);
129        Ok(Self { names, lines })
130    }
131}