Skip to main content

wowlab_cli/commands/snapshot/db/
connection.rs

1use 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
48/// Compare-and-swap `game.meta` to `target` revision, returning rows affected.
49pub(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}