path: root/rust/scraper/src/database/postgres.rs

use std::mem;
use std::time::Duration;
use crate::database::error::DatabaseError;
use chrono::{DateTime, NaiveDate, TimeZone, Utc};
use proto::touchpad::common::v1::{Event, EventTime, EventTimeResult, Gender, individual_event, IndividualEvent, relay_event, RelayEvent, Stroke, SwimMeet, Swimmer, Team};
use sqlx::postgres::{PgPool, PgPoolOptions, PgQueryResult};
use proto::ProtoTimestamp;
use proto::touchpad::common::v1;

pub struct PostgresClient {
    pool: PgPool,
}

impl PostgresClient {
    pub async fn new<S: AsRef<str>>(url: S) -> Result<PostgresClient, sqlx::Error> {
        Ok(PostgresClient {
            pool: PgPoolOptions::new().connect(url.as_ref()).await?,
        })
    }
}

#[async_trait::async_trait]
impl super::DatabaseClient for PostgresClient {
    async fn get_swimmer(&self, id: u32) -> Result<Swimmer, DatabaseError> {
        let swimmer: SwimmersTableSchema = sqlx::query_as!(
            SwimmersTableSchema,
            "SELECT * FROM swimmers WHERE id = $1;",
            id as i32
        )
            .fetch_one(&self.pool)
            .await?;

        Ok(Swimmer {
            id,
            name: swimmer.name,
            gender: Gender::from_iso_5218(swimmer.gender as u8) as i32,
        })
    }

    async fn add_swimmer(&self, swimmer: &Swimmer) -> Result<(), DatabaseError> {
        let gender = Gender::from_i32(swimmer.gender)
            .ok_or(DatabaseError::BadInput("Invalid gender".into()))?
            .to_iso_5218();

        sqlx::query!(
            "INSERT INTO swimmers VALUES ($1, $2, $3);",
            swimmer.id as i32,
            swimmer.name,
            gender as i16
        )
            .execute(&self.pool)
            .await?;

        Ok(())
    }

    async fn get_team(&self, id: u32) -> crate::database::Result<Team> {
        let team: TeamsTableSchema = sqlx::query_as!(
            TeamsTableSchema,
            "SELECT * FROM teams WHERE id = $1;",
            id as i32
        )
            .fetch_one(&self.pool)
            .await?;

        Ok(Team {
            id,
            name: team.name,
        })
    }

    async fn add_team(&self, team: &Team) -> crate::database::Result<()> {
        sqlx::query!(
            "INSERT INTO teams VALUES ($1, $2);",
            team.id as i32,
            team.name
        )
            .execute(&self.pool)
            .await?;

        Ok(())
    }

    async fn get_meet(&self, id: u32) -> crate::database::Result<SwimMeet> {
        let mut meet: Vec<GetMeetSchema> = sqlx::query_as!(
            GetMeetSchema,
            "
SELECT meets.name as meet_name, start_date, end_date, team_id, t.name as team_name
FROM meets
INNER JOIN meets_teams mt on meets.id = mt.meet_id
INNER JOIN teams t on mt.team_id = t.id
WHERE meets.id = $1;
",
            id as i32
        )
            .fetch_all(&self.pool)
            .await?;

        let meet_name = mem::take(&mut meet.get_mut(0).ok_or(DatabaseError::NotFound)?.meet_name);
        let start = meet.get(0).ok_or(DatabaseError::NotFound)?.start_date.and_hms(0, 0, 0);
        let end = meet.get(0).ok_or(DatabaseError::NotFound)?.end_date.and_hms(0, 0, 0);

        let teams = meet.into_iter()
            .map(|m| Team {
                id: m.team_id as u32,
                name: m.team_name,
            })
            .collect();

        Ok(SwimMeet {
            id,
            meet_name,
            start: Some(Utc.from_utc_datetime(&start).into()),
            end: Some(Utc.from_utc_datetime(&end).into()),
            teams,
            points: Default::default(),
        })
    }

    async fn add_meet(&self, meet: &SwimMeet) -> crate::database::Result<()> {
        let mut tx = self.pool.begin().await?;

        // Convert protobuf timestamp into chrono Datetime
        let start_date: DateTime<Utc> = {
            let proto_timestamp: ProtoTimestamp = meet.start.clone()
                .ok_or(DatabaseError::BadInput("empty start date".into()))?
                .into();
            proto_timestamp.into()
        };

        let end_date: DateTime<Utc> = {
            let proto_timestamp: ProtoTimestamp = meet.end.clone()
                .ok_or(DatabaseError::BadInput("empty end date".into()))?
                .into();
            proto_timestamp.into()
        };

        // Insert into meets table
        sqlx::query!(
            "INSERT INTO meets VALUES ($1, $2, $3, $4)",
            meet.id as i32,
            meet.meet_name,
            start_date.date_naive(),
            end_date.date_naive(),
        )
            .execute(&mut tx)
            .await?;

        for team in &meet.teams {
            sqlx::query!(
                "INSERT INTO teams VALUES ($1, $2) ON CONFLICT DO NOTHING",
                team.id as i32,
                team.name,
            )
                .execute(&mut tx)
                .await?;

            sqlx::query!(
                "INSERT INTO meets_teams VALUES ($1, $2)",
                meet.id as i32,
                team.id as i32,
            )
                .execute(&mut tx)
                .await?;
        }

        tx.commit().await?;

        Ok(())
    }

    async fn get_event(&self, id: u32) -> crate::database::Result<Event> {
        let event: Vec<GetEventSchema> = sqlx::query_as!(
            GetEventSchema,
"
SELECT *
FROM events e
LEFT JOIN timings t on e.id = t.event_id
WHERE e.id = $1;
",
            id as i32,
        )
            .fetch_all(&self.pool)
            .await?;

        db_event_to_proto(event)
    }

    async fn get_event_by_number(&self, meet_id: u32, number: u32) -> crate::database::Result<Event> {
        let event: Vec<GetEventSchema> = sqlx::query_as!(
            GetEventSchema,
"
SELECT *
FROM events e
LEFT JOIN timings t on e.id = t.event_id
WHERE e.meet_id = $1 AND e.event_num = $2;
",
            meet_id as i32,
            number as i32,
        )
            .fetch_all(&self.pool)
            .await?;

        db_event_to_proto(event)
    }

    async fn add_event(&self, event: &Event) -> crate::database::Result<()> {
        let mut tx = self.pool.begin().await?;

        let stroke = Stroke::from_i32(event.stroke)
            .ok_or(DatabaseError::BadInput("invalid stroke".into()))?;
        let gender = Gender::from_i32(event.gender)
            .ok_or(DatabaseError::BadInput("invalid gender".into()))?
            .to_iso_5218();

        sqlx::query!(
            "INSERT INTO events VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9);",
            event.id as i32,
            event.meet_id as i32,
            event.event_num as i16,
            event.age_lo as i16,
            event.age_hi as i16,
            event.distance as i32,
            stroke as i16,
            gender as i16,
            event.is_relay
        )
            .execute(&mut tx)
            .await?;

        match &event.event {
            Some(v1::event::Event::Individual(ev)) => {
                for t in &ev.times {
                    if let Some(time) = &t.time {
                        insert_event_timing(
                            &mut tx,
                            time,
                            None,
                            Some(t.swimmer_id),
                        ).await?;
                    }
                }
            }
            Some(v1::event::Event::Relay(ev)) => {
                for t in &ev.times {
                    if let Some(time) = &t.time {
                        insert_event_timing(
                            &mut tx,
                            time,
                            Some(&t.name),
                            None,
                        ).await?;
                    }
                }
            }
            None => {}
        }

        tx.commit().await?;

        Ok(())
    }
}

struct SwimmersTableSchema {
    id: i32,
    name: String,
    gender: i16,
}

struct TeamsTableSchema {
    id: i32,
    name: String,
}

struct GetMeetSchema {
    meet_name: String,
    start_date: NaiveDate,
    end_date: NaiveDate,
    team_id: i32,
    team_name: String,
}

struct GetEventSchema {
    id: i32,
    meet_id: i32,
    event_num: i16,
    age_lo: i16,
    age_hi: i16,
    distance: i32,
    stroke: i16,
    gender: i16,
    is_relay: bool,

    name: Option<String>,
    swimmer_id: Option<i32>,
    event_id: i32,
    team_id: i32,
    heat: i16,
    lane: i16,
    result: i16,
    final_time: Option<i64>,
    seed_time: Option<i64>,
    rank: i16,
    points: f32,
}

fn db_event_to_proto(event: Vec<GetEventSchema>) -> crate::database::Result<Event> {
    let first = event.first().ok_or(DatabaseError::NotFound)?;

    let mut proto_event = Event {
        id: first.event_id as u32,
        meet_id: first.meet_id as u32,
        event_num: first.event_num as u32,
        age_lo: first.age_lo as u32,
        age_hi: first.age_hi as u32,
        distance: first.distance as u32,
        stroke: Stroke::from_i32(first.stroke as i32).unwrap_or_default() as i32,
        gender: Gender::from_iso_5218(first.stroke as u8) as i32,
        is_relay: first.is_relay,
        event: None,
    };

    if proto_event.is_relay {
        proto_event.event = Some(v1::event::Event::Relay(RelayEvent {
            times: event.into_iter()
                .map(|mut e| relay_event::RelayEventTime {
                    name: e.name.take().unwrap_or_default(),
                    time: Some(create_proto_ev_time(&e)),
                })
                .collect()
        }));
    } else {
        proto_event.event = Some(v1::event::Event::Individual(IndividualEvent {
            times: event.into_iter()
                .map(|e| individual_event::IndividualEventTime {
                    swimmer_id: e.swimmer_id.unwrap_or_default() as u32,
                    time: Some(create_proto_ev_time(&e)),
                })
                .collect()
        }))
    }

    Ok(proto_event)
}

fn create_proto_ev_time(e: &GetEventSchema) -> EventTime {
    EventTime {
        event_id: e.event_id as u32,
        team_id: e.team_id as u32,
        heat: e.heat as u32,
        lane: e.lane as u32,
        result: EventTimeResult::from_i32(e.result as i32).unwrap_or_default() as i32,
        time: e.final_time.map(|t| Duration::from_millis(t as u64).into()),
        seed: e.seed_time.map(|t| Duration::from_millis(t as u64).into()),
        rank: e.rank as i32,
        points: e.points,
    }
}

async fn insert_event_timing(
    ex: impl sqlx::PgExecutor<'_>,
    time: &EventTime,
    name: Option<&str>,
    swimmer_id: Option<u32>,
) -> Result<PgQueryResult, sqlx::Error> {
    let result = EventTimeResult::from_i32(time.result)
        .unwrap_or_default();

    let final_time: Option<i64> = time.time.clone()
        .and_then(|t| t.try_into().ok())
        .and_then(|t: Duration| t.as_millis().try_into().ok());
    let seed_time: Option<i64> = time.seed.clone()
        .and_then(|t| t.try_into().ok())
        .and_then(|t: Duration| t.as_millis().try_into().ok());

    sqlx::query!(
        "INSERT INTO timings VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11);",
        name,
        swimmer_id.map(|id| id as i32),
        time.event_id as i32,
        time.team_id as i32,
        time.heat as i16,
        time.lane as i16,
        result as i16,
        time.rank as i32,
        time.points,
        final_time,
        seed_time,
    )
        .execute(ex)
        .await
}