septastic/api/src/database.rs

use std::collections::HashMap;
use sqlx::{Postgres, Transaction};
use libseptastic::{stop_schedule::{Trip, TripTracking, StopSchedule}};

pub async fn get_route_by_id(
    id: String,
    transaction: &mut Transaction<'_, Postgres>,
) -> ::anyhow::Result<libseptastic::route::Route> {

    let row = sqlx::query!(
        r#"SELECT 
            id, 
            name, 
            short_name, 
            color_hex, 
            route_type as "route_type: libseptastic::route::RouteType"
        FROM 
            septa_routes
        WHERE 
            id = $1
        ;"#,
        id
    )
    .fetch_one(&mut **transaction)
    .await?;

    return Ok(libseptastic::route::Route {
        name: row.name,
        short_name: row.short_name,
        color_hex: row.color_hex,
        route_type: row.route_type,
        id: row.id,
    });
}

pub async fn get_all_routes(
    transaction: &mut Transaction<'_, Postgres>,
) -> ::anyhow::Result<Vec<libseptastic::route::Route>> {

    let rows = sqlx::query!(
        r#"SELECT 
            id, 
            name, 
            short_name, 
            color_hex, 
            route_type as "route_type: libseptastic::route::RouteType"
        FROM 
            septa_routes
        ORDER BY
            CASE
                WHEN id ~ '^[0-9]+$' THEN CAST(id AS INT)
                ELSE NULL
            END ASC,
            id ASC;
        ;"#
    )
    .fetch_all(&mut **transaction)
    .await?;

    let mut routes = Vec::new();

    for row in rows {
        routes.push(libseptastic::route::Route {
        name: row.name,
        short_name: row.short_name,
        color_hex: row.color_hex,
        route_type: row.route_type,
        id: row.id,
        });
    }

    return Ok(routes);
}

pub async fn get_direction_by_route_id(
    id: String,
    transaction: &mut Transaction<'_, Postgres>,
) -> ::anyhow::Result<Vec<libseptastic::direction::Direction>> {

    let rows = sqlx::query!(
        r#"SELECT 
            route_id,
            direction_id,
            direction as "direction: libseptastic::direction::CardinalDirection",
            direction_destination
        FROM 
            septa_directions
        WHERE 
            route_id = $1
        ;"#,
        id
    )
    .fetch_all(&mut **transaction)
    .await?;

    let mut res = Vec::new();

    for row in rows {
        res.push(libseptastic::direction::Direction{
            route_id: row.route_id,
            direction_id: row.direction_id,
            direction: row.direction,
            direction_destination: row.direction_destination
        });
    }

    return Ok(res);
}

pub async fn get_schedule_by_route_id(
    id: String,
    transaction: &mut Transaction<'_, Postgres>,
) -> ::anyhow::Result<Vec<Trip>> {

    let schedule_day = chrono::Utc::now().with_timezone(&chrono_tz::America::New_York);
    let schedule_day_str = schedule_day.format("%Y%m%d").to_string();

    let rows = sqlx::query!(
        r#"SELECT 
            septa_stop_schedules.route_id,
            septa_stops.name as stop_name,
            trip_id, 
            septa_stop_schedules.service_id, 
            septa_stop_schedules.direction_id,
            arrival_time,
            stop_id,
            stop_sequence
        FROM 
            septa_stop_schedules
        INNER JOIN septa_stops 
            ON septa_stops.id = septa_stop_schedules.stop_id
        INNER JOIN septa_schedule_days
            ON septa_schedule_days.date = $2 
                AND
               septa_schedule_days.service_id = septa_stop_schedules.service_id
        WHERE 
            septa_stop_schedules.route_id = $1
        ;"#,
        id.clone(),
        schedule_day_str.clone()
    )
    .fetch_all(&mut **transaction)
    .await?;

    let mut sched_groups: HashMap<String, Vec<StopSchedule>> = HashMap::new();
    for row in rows {
        let arr = match sched_groups.get_mut(&row.trip_id) {
            Some(x) => x,
            None => {
                sched_groups.insert(row.trip_id.clone(), Vec::new());
                sched_groups.get_mut(&row.trip_id).unwrap()
            }
        };

        arr.push(StopSchedule {
            route_id: row.route_id,
            stop_name: row.stop_name,
            trip_id: row.trip_id,
            service_id: row.service_id,
            direction_id: row.direction_id,
            arrival_time: row.arrival_time,
            stop_id: row.stop_id,
            stop_sequence: row.stop_sequence
        });
    }
    let mut res = Vec::new();

    for group in sched_groups {
        res.push(Trip{
            trip_id: group.0,
            route_id: group.1[0].route_id.clone(),
            schedule: group.1.clone(),
            direction_id: group.1[0].direction_id.clone(),
            tracking_data: TripTracking::Untracked 
        });
    }

    return Ok(res);
}