team archive

archiver/src/archive.rs

@@ -1,8 +1,13 @@
+use futures_util::stream::TryStreamExt;
 use sqlx::postgres::PgConnectOptions;
-use sqlx::{query, Error, PgPool};
+use sqlx::{query, Error, Executor, PgPool, Postgres};
+use std::ops::Range;
 use std::str::FromStr;
 use thiserror::Error;
-use ugc_scraper_types::MatchInfo;
+use tokio_stream::Stream;
+use ugc_scraper_types::{
+    GameMode, MatchInfo, Membership, MembershipRole, NameChange, Record, Region, Team,
+};
 
 #[derive(Debug, Error)]
 pub enum ArchiveError {
@@ -70,4 +75,197 @@ impl Archive {
             })?
             .map(|row| row.id as u32))
     }
+
+    pub async fn store_team(&self, id: u32, team: &Team) -> Result<(), ArchiveError> {
+        let mut transaction = self
+            .pool
+            .begin()
+            .await
+            .map_err(|error| ArchiveError::Query {
+                description: "beginning team transaction",
+                error,
+            })?;
+        query!(
+            "INSERT INTO teams (
+                id, tag, name, image, format, region, timezone, steam_group, division, description
+              ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)",
+            id as i32,
+            team.tag,
+            team.name,
+            team.image,
+            team.format as GameMode,
+            team.region as Option<Region>,
+            team.timezone,
+            team.steam_group,
+            team.division,
+            team.description,
+        )
+        .execute(&mut *transaction)
+        .await
+        .map_err(|error| ArchiveError::Query {
+            description: "inserting team",
+            error,
+        })?;
+
+        for title in team.titles.iter() {
+            Self::store_title(&mut *transaction, id, title).await?;
+        }
+        for name_change in team.name_changes.iter() {
+            Self::store_team_name_change(&mut *transaction, id, name_change).await?
+        }
+        for record in team.results.iter() {
+            Self::store_record(&mut *transaction, id, record).await?
+        }
+        for membership in team.members.iter() {
+            Self::store_membership(&mut *transaction, id, membership).await?
+        }
+
+        transaction
+            .commit()
+            .await
+            .map_err(|error| ArchiveError::Query {
+                description: "commiting team transaction",
+                error,
+            })?;
+        Ok(())
+    }
+
+    pub async fn update_team_region(&self, id: u32, team: &Team) -> Result<(), ArchiveError> {
+        query!(
+            "UPDATE teams SET region = $2 WHERE id = $1",
+            id as i32,
+            team.region as Option<Region>,
+        )
+        .execute(&self.pool)
+        .await
+        .map_err(|error| ArchiveError::Query {
+            description: "updating team region",
+            error,
+        })?;
+        Ok(())
+    }
+
+    pub async fn get_last_team_id(&self) -> Result<Option<u32>, ArchiveError> {
+        Ok(query!("SELECT id FROM teams ORDER BY id DESC LIMIT 1")
+            .fetch_optional(&self.pool)
+            .await
+            .map_err(|error| ArchiveError::Query {
+                description: "getting latest team",
+                error,
+            })?
+            .map(|row| row.id as u32))
+    }
+
+    pub async fn get_team_range(&self) -> Result<Range<u32>, ArchiveError> {
+        let row = query!("select greatest(max(team_home), max(team_away)) as max, least(min(team_home), min(team_away)) as min from matches limit 1;")
+            .fetch_one(&self.pool)
+            .await
+            .map_err(|error| ArchiveError::Query {
+                description: "getting latest team",
+                error,
+            })?;
+        Ok((row.min.unwrap_or_default() as u32)..(row.max.unwrap_or_default() as u32))
+    }
+
+    pub fn get_no_region_teams(&self) -> impl Stream<Item = Result<u32, ArchiveError>> + use<'_> {
+        query!("select id from teams where region IS NULL and format != 'eights' order by id desc")
+            .fetch(&self.pool)
+            .map_err(|error| ArchiveError::Query {
+                description: "getting teams without region",
+                error,
+            })
+            .map_ok(|map| map.id as u32)
+    }
+
+    async fn store_title(
+        db: impl Executor<'_, Database = Postgres>,
+        team_id: u32,
+        title: &str,
+    ) -> Result<(), ArchiveError> {
+        query!(
+            "INSERT INTO titles (
+                team_id, title
+              ) VALUES ($1, $2)",
+            team_id as i32,
+            title
+        )
+        .execute(db)
+        .await
+        .map_err(|error| ArchiveError::Query {
+            description: "inserting title",
+            error,
+        })?;
+        Ok(())
+    }
+
+    async fn store_team_name_change(
+        db: impl Executor<'_, Database = Postgres>,
+        team_id: u32,
+        change: &NameChange,
+    ) -> Result<(), ArchiveError> {
+        query!(
+            "INSERT INTO team_name_changes (
+                team_id, from_tag, from_name, to_tag, to_name, date
+              ) VALUES ($1, $2, $3, $4, $5, $6)",
+            team_id as i32,
+            change.from_tag,
+            change.from,
+            change.to_tag,
+            change.to,
+            change.date
+        )
+        .execute(db)
+        .await
+        .map_err(|error| ArchiveError::Query {
+            description: "inserting name change",
+            error,
+        })?;
+        Ok(())
+    }
+
+    async fn store_membership(
+        db: impl Executor<'_, Database = Postgres>,
+        team_id: u32,
+        membership: &Membership,
+    ) -> Result<(), ArchiveError> {
+        query!(
+            "INSERT INTO memberships (
+                team_id, steam_id, role, since
+              ) VALUES ($1, $2, $3, $4)",
+            team_id as i32,
+            u64::from(membership.steam_id) as i64,
+            membership.role as MembershipRole,
+            membership.since,
+        )
+        .execute(db)
+        .await
+        .map_err(|error| ArchiveError::Query {
+            description: "inserting membership",
+            error,
+        })?;
+        Ok(())
+    }
+
+    async fn store_record(
+        db: impl Executor<'_, Database = Postgres>,
+        team_id: u32,
+        record: &Record,
+    ) -> Result<(), ArchiveError> {
+        query!(
+            "INSERT INTO records (
+                team_id, season, wins, losses
+              ) VALUES ($1, $2, $3, $4)",
+            team_id as i32,
+            record.season as i32,
+            record.wins as i32,
+            record.losses as i32,
+        )
+        .execute(db)
+        .await
+        .map_err(|error| ArchiveError::Query {
+            description: "inserting record",
+            error,
+        })?;
+        Ok(())
+    }
 }

archiver/src/main.rs

@@ -8,8 +8,10 @@ use crate::config::Config;
 use clap::{Parser, Subcommand};
 use main_error::MainResult;
 use std::path::PathBuf;
+use std::pin::pin;
 use std::time::Duration;
 use tokio::time::sleep;
+use tokio_stream::StreamExt;
 use tracing::{error, info, span, warn, Level};
 
 #[derive(Debug, Parser)]
@@ -23,6 +25,8 @@ struct Args {
 #[derive(Debug, Subcommand)]
 enum Command {
     Matches,
+    Teams,
+    FixupTeams,
 }
 
 const LAST_MATCH: u32 = 117047;
@@ -40,6 +44,12 @@ async fn main() -> MainResult {
         Command::Matches => {
             archive_matches(&client, &archive).await?;
         }
+        Command::Teams => {
+            archive_teams(&client, &archive).await?;
+        }
+        Command::FixupTeams => {
+            fixup_teams(&client, &archive).await?;
+        }
     }
     Ok(())
 }
@@ -69,6 +79,61 @@ async fn archive_matches(client: &UgcClient, archive: &Archive) -> MainResult {
     Ok(())
 }
 
+async fn archive_teams(client: &UgcClient, archive: &Archive) -> MainResult {
+    let range = archive.get_team_range().await?;
+    let next_team = archive.get_last_team_id().await?.unwrap_or(range.start - 1) + 1;
+
+    for id in next_team..=range.end {
+        let _span = span!(Level::INFO, "archive_team", id = id).entered();
+        match client.get_team(id).await.check_not_found() {
+            Ok(Some(team_data)) => {
+                if team_data.format.is_tf2() {
+                    info!("storing team");
+                    archive.store_team(id, &team_data).await?;
+                } else {
+                    info!("skipping non-tf2 team");
+                }
+            }
+            Ok(None) => {
+                warn!("team not found");
+            }
+            Err(e) => {
+                error!("error fetching team: {:?}", e);
+                panic!();
+            }
+        }
+        sleep(Duration::from_millis(500)).await;
+    }
+    Ok(())
+}
+
+async fn fixup_teams(client: &UgcClient, archive: &Archive) -> MainResult {
+    let mut ids = pin!(archive.get_no_region_teams());
+
+    while let Some(Ok(id)) = ids.next().await {
+        let _span = span!(Level::INFO, "fixup_team", id = id).entered();
+        match client.get_team(id).await.check_not_found() {
+            Ok(Some(team_data)) => {
+                if team_data.format.is_tf2() {
+                    info!(region = ?team_data.region, "updating team region");
+                    archive.update_team_region(id, &team_data).await?;
+                } else {
+                    info!("skipping non-tf2 team");
+                }
+            }
+            Ok(None) => {
+                warn!("team not found");
+            }
+            Err(e) => {
+                error!("error fetching team: {:?}", e);
+                panic!();
+            }
+        }
+        sleep(Duration::from_millis(500)).await;
+    }
+    Ok(())
+}
+
 trait NotFoundResultExt<T>: Sized {
     fn check_not_found(self) -> Result<Option<T>, UgcClientError>;
 }