1use crate::Lb;
2use crate::model::api::{FileUsage, GetUsageRequest};
3use crate::model::errors::LbResult;
4use crate::model::file_like::FileLike;
5use crate::model::tree_like::TreeLike;
6use crate::model::usage::{bytes_to_human, get_usage};
7use serde::Serialize;
8use std::collections::HashMap;
9use uuid::Uuid;
10
11#[derive(Serialize, Debug, Clone)]
12pub struct UsageMetrics {
13 pub usages: Vec<FileUsage>,
14 pub server_usage: UsageItemMetric,
15 pub data_cap: UsageItemMetric,
16}
17
18#[derive(Serialize, PartialEq, Eq, Debug, Clone)]
19pub struct UsageItemMetric {
20 pub exact: u64,
21 pub readable: String,
22}
23
24impl Lb {
25 #[instrument(level = "debug", skip(self))]
33 pub async fn get_usage(&self) -> LbResult<UsageMetrics> {
34 let acc = self.get_account()?;
35 let usage = self.client.request(acc, GetUsageRequest {}).await?;
36 Ok(get_usage(usage))
37 }
38
39 #[instrument(level = "debug", skip(self))]
40 pub async fn get_uncompressed_usage_breakdown(&self) -> LbResult<HashMap<Uuid, usize>> {
41 let tx = self.ro_tx().await;
42 let db = tx.db();
43
44 let mut tree = (&db.base_metadata).to_staged(&db.local_metadata).to_lazy();
45 let mut breakdown = HashMap::default();
46
47 for id in tree.ids() {
48 let is_file_deleted = tree.calculate_deleted(&id)?;
49 let file = tree.find(&id)?;
50
51 if !is_file_deleted && file.is_document() {
52 let doc = self.read_document_helper(id, &mut tree).await?;
53 let doc_size = doc.len();
54 breakdown.insert(id, doc_size);
55 }
56 }
57
58 Ok(breakdown)
59 }
60
61 #[instrument(level = "debug", skip(self))]
63 pub async fn get_uncompressed_usage(&self) -> LbResult<UsageItemMetric> {
64 let tx = self.ro_tx().await;
65 let db = tx.db();
66
67 let mut tree = (&db.base_metadata).to_staged(&db.local_metadata).to_lazy();
68
69 let mut local_usage: u64 = 0;
70 for id in tree.ids() {
71 let is_file_deleted = tree.calculate_deleted(&id)?;
72 let file = tree.find(&id)?;
73
74 if !is_file_deleted && file.is_document() {
75 let doc = self.read_document_helper(id, &mut tree).await?;
76 local_usage += doc.len() as u64
77 }
78 }
79
80 let readable = bytes_to_human(local_usage);
81 Ok(UsageItemMetric { exact: local_usage, readable })
82 }
83}