wowlab_cli/commands/snapshot/db/
connection.rs1use sqlx::{Executor, PgPool, postgres::PgPoolOptions};
2use wowlab_types::{sim::FastMap, table_registry::GameDataTable};
3
4pub(crate) async fn connect(url: &str) -> Result<PgPool, sqlx::Error> {
5 PgPoolOptions::new()
6 .after_connect(|conn, _meta| {
7 Box::pin(async move {
8 conn.execute("SET statement_timeout = 0").await?;
9 conn.execute("SET synchronous_commit = off").await?;
10
11 Ok(())
12 })
13 })
14 .connect(url)
15 .await
16}
17
18pub(crate) fn should_sync(tables: &[GameDataTable], table: GameDataTable) -> bool {
19 tables.is_empty() || tables.contains(&table)
20}
21
22pub(crate) async fn count_rows(
23 pool: &PgPool,
24 tables: &[&str],
25) -> Result<FastMap<String, i64>, sqlx::Error> {
26 let mut counts = FastMap::default();
27
28 counts.reserve(tables.len());
29
30 for &table in tables {
31 let query = format!("SELECT COUNT(*) as count FROM {table}");
32 let row: (i64,) = sqlx::query_as(&query).fetch_one(pool).await?;
33
34 counts.insert(table.to_string(), row.0);
35 }
36
37 Ok(counts)
38}
39
40pub(crate) async fn read_meta_revision(pool: &PgPool) -> Result<i64, sqlx::Error> {
41 let row: Option<(i64,)> = sqlx::query_as("SELECT revision FROM game.meta WHERE id = 1")
42 .fetch_optional(pool)
43 .await?;
44
45 Ok(row.map_or(0, |r| r.0))
46}
47
48pub(crate) async fn write_meta(
50 pool: &PgPool,
51 patch: &str,
52 expected_current: i64,
53 target: i64,
54 table_revisions: serde_json::Value,
55) -> Result<u64, sqlx::Error> {
56 let result = sqlx::query(
57 "INSERT INTO game.meta (id, revision, patch_version, table_revisions)
58 VALUES (1, $3, $1, $4)
59 ON CONFLICT (id) DO UPDATE SET
60 revision = $3,
61 patch_version = $1,
62 table_revisions = $4,
63 synced_at = now()
64 WHERE game.meta.revision = $2",
65 )
66 .bind(patch)
67 .bind(expected_current)
68 .bind(target)
69 .bind(table_revisions)
70 .execute(pool)
71 .await?;
72
73 Ok(result.rows_affected())
74}