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}