wowlab_cli/commands/snapshot/db/
inserts.rs

1use sqlx::PgPool;
2use wowlab_types::{
3    copy::{CopyRow, to_json},
4    data::{MythicPlusSeasonFlat, SpecDataFlat},
5};
6
7use super::copy::{COPY_FLUSH_BYTES, progress_bar, upsert_all_columns};
8
9const ASYNC_YIELD_ROWS: usize = 4_096;
10
11pub(crate) async fn insert_specs(
12    pool: &PgPool,
13    rows: &[SpecDataFlat],
14    patch: &str,
15) -> Result<(), sqlx::Error> {
16    const COLUMNS: &[&str] = &[
17        "id",
18        "patch_version",
19        "name",
20        "description",
21        "class_id",
22        "class_name",
23        "role",
24        "order_index",
25        "icon_file_id",
26        "file_name",
27        "primary_stat_priority",
28        "mastery_spell_id_0",
29        "mastery_spell_id_1",
30    ];
31
32    let pb = progress_bar(rows.len(), "Specs");
33    let mut tx = pool.begin().await?;
34
35    sqlx::query(
36        "CREATE TEMP TABLE specs_staging (LIKE game.specs INCLUDING DEFAULTS) ON COMMIT DROP",
37    )
38    .execute(&mut *tx)
39    .await?;
40
41    let stmt = format!("COPY specs_staging ({}) FROM STDIN", COLUMNS.join(", "));
42    let mut sink = tx.copy_in_raw(&stmt).await?;
43    let mut writer = CopyRow::new();
44
45    for (index, spec) in rows.iter().enumerate() {
46        writer
47            .push_bind(spec.id)
48            .push_bind(patch)
49            .push_bind(spec.name.as_str())
50            .push_bind(&spec.description)
51            .push_bind(spec.class_id)
52            .push_bind(spec.class_name.as_str())
53            .push_bind(spec.role)
54            .push_bind(spec.order_index)
55            .push_bind(spec.icon_file_id)
56            .push_bind(spec.file_name.as_str())
57            .push_bind(spec.primary_stat_priority)
58            .push_bind(spec.mastery_spell_id_0)
59            .push_bind(spec.mastery_spell_id_1);
60        writer.finish_row();
61
62        if index % ASYNC_YIELD_ROWS == 0 {
63            tokio::task::yield_now().await;
64        }
65    }
66
67    sink.send(writer.take()).await?;
68    sink.finish().await?;
69
70    let cols = COLUMNS.join(", ");
71    let merge = format!(
72        "INSERT INTO game.specs ({cols}) SELECT {cols} FROM specs_staging \
73         ON CONFLICT (id) DO UPDATE SET {}",
74        upsert_all_columns(COLUMNS)
75    );
76
77    sqlx::query(&merge).execute(&mut *tx).await?;
78    sqlx::query("DELETE FROM game.specs WHERE id NOT IN (SELECT id FROM specs_staging)")
79        .execute(&mut *tx)
80        .await?;
81    tx.commit().await?;
82    pb.finish("");
83
84    Ok(())
85}
86
87pub(crate) async fn insert_mythic_plus_seasons(
88    pool: &PgPool,
89    rows: &[MythicPlusSeasonFlat],
90    patch: &str,
91) -> Result<(), sqlx::Error> {
92    const COLUMNS: &[&str] = &[
93        "id",
94        "patch_version",
95        "key_rewards",
96        "milestone_season",
97        "start_time_event",
98        "expansion_level",
99        "heroic_lfg_dungeon_min_gear",
100        "display_season_id",
101        "display_season_name",
102        "display_season_index",
103        "delves_season_id",
104        "item_group_ilvl_scaling_id",
105        "bonus_list_groups",
106        "crest_currencies",
107        "valorstones_currency_id",
108        "tracked_dungeons",
109        "tracked_affixes",
110        "key_floors",
111    ];
112
113    let pb = progress_bar(rows.len(), "MythicPlusSeasons");
114    let mut tx = pool.begin().await?;
115
116    sqlx::query("TRUNCATE game.mythic_plus_seasons RESTART IDENTITY")
117        .execute(&mut *tx)
118        .await?;
119
120    let stmt = format!(
121        "COPY game.mythic_plus_seasons ({}) FROM STDIN",
122        COLUMNS.join(", ")
123    );
124    let mut sink = tx.copy_in_raw(&stmt).await?;
125    let mut writer = CopyRow::new();
126
127    for (index, row) in rows.iter().enumerate() {
128        writer
129            .push_bind(row.id)
130            .push_bind(patch)
131            .push_bind(to_json(&row.key_rewards))
132            .push_bind(row.milestone_season)
133            .push_bind(row.start_time_event)
134            .push_bind(row.expansion_level)
135            .push_bind(row.heroic_lfg_dungeon_min_gear)
136            .push_bind(row.display_season_id)
137            .push_bind(&row.display_season_name)
138            .push_bind(row.display_season_index)
139            .push_bind(row.delves_season_id)
140            .push_bind(row.item_group_ilvl_scaling_id)
141            .push_bind(&row.bonus_list_groups)
142            .push_bind(&row.crest_currencies)
143            .push_bind(row.valorstones_currency_id)
144            .push_bind(to_json(&row.tracked_dungeons))
145            .push_bind(to_json(&row.tracked_affixes))
146            .push_bind(to_json(&row.key_floors));
147        writer.finish_row();
148
149        if writer.len() >= COPY_FLUSH_BYTES {
150            sink.send(writer.take()).await?;
151            pb.tick(index as u64 + 1);
152        } else if index % ASYNC_YIELD_ROWS == 0 {
153            tokio::task::yield_now().await;
154        }
155    }
156
157    sink.send(writer.take()).await?;
158    sink.finish().await?;
159    tx.commit().await?;
160    pb.finish("");
161
162    Ok(())
163}
wowlab_cli/commands/snapshot/db/inserts.rs

wowlab_cli/commands/snapshot/db/
inserts.rs
