usages.rs

 1use std::str::FromStr;
 2use strum::IntoEnumIterator as _;
 3
 4use super::*;
 5
 6impl LlmDatabase {
 7    pub async fn initialize_usage_measures(&mut self) -> Result<()> {
 8        let all_measures = self
 9            .transaction(|tx| async move {
10                let existing_measures = usage_measure::Entity::find().all(&*tx).await?;
11
12                let new_measures = UsageMeasure::iter()
13                    .filter(|measure| {
14                        !existing_measures
15                            .iter()
16                            .any(|m| m.name == measure.to_string())
17                    })
18                    .map(|measure| usage_measure::ActiveModel {
19                        name: ActiveValue::set(measure.to_string()),
20                        ..Default::default()
21                    })
22                    .collect::<Vec<_>>();
23
24                if !new_measures.is_empty() {
25                    usage_measure::Entity::insert_many(new_measures)
26                        .exec(&*tx)
27                        .await?;
28                }
29
30                Ok(usage_measure::Entity::find().all(&*tx).await?)
31            })
32            .await?;
33
34        self.usage_measure_ids = all_measures
35            .into_iter()
36            .filter_map(|measure| {
37                UsageMeasure::from_str(&measure.name)
38                    .ok()
39                    .map(|um| (um, measure.id))
40            })
41            .collect();
42        Ok(())
43    }
44}