save_downloads.rs

use anyhow::Context;
use chrono::NaiveDate;
use crates_io_cdn_logs::DownloadsMap;
use diesel::prelude::*;
use diesel::PgConnection;

table! {
    temp_downloads (name, version, date) {
        name -> Text,
        version -> Text,
        date -> Date,
        downloads -> BigInt,
    }
}

#[derive(Insertable)]
#[diesel(table_name = temp_downloads)]
struct NewDownload {
    name: String,
    version: String,
    date: NaiveDate,
    downloads: i64,
}

pub fn save_downloads(downloads: DownloadsMap, conn: &mut PgConnection) -> anyhow::Result<()> {
    conn.transaction(|conn| {
        debug!("Creating temp_downloads table");
        create_temp_downloads_table(conn).context("Failed to create temp_downloads table")?;

        debug!("Saving counted downloads to temp_downloads table");
        fill_temp_downloads_table(downloads, conn)
            .context("Failed to fill temp_downloads table")?;

        debug!("Saving temp_downloads to version_downloads table");
        save_to_version_downloads(conn)
            .context("Failed to save temp_downloads to version_downloads table")?;

        Ok(())
    })
}

fn create_temp_downloads_table(conn: &mut PgConnection) -> QueryResult<usize> {
    diesel::sql_query(
        r#"
            CREATE TEMPORARY TABLE temp_downloads (
                name VARCHAR NOT NULL,
                version VARCHAR NOT NULL,
                date DATE NOT NULL,
                downloads INTEGER NOT NULL,
                PRIMARY KEY (name, version, date)
            ) ON COMMIT DROP;
        "#,
    )
    .execute(conn)
}

fn fill_temp_downloads_table(
    downloads: DownloadsMap,
    conn: &mut PgConnection,
) -> QueryResult<usize> {
    let map = downloads
        .into_inner()
        .into_iter()
        .map(|((name, version, date), downloads)| NewDownload {
            name,
            version: version.to_string(),
            date,
            downloads: downloads as i64,
        })
        .collect::<Vec<_>>();

    diesel::insert_into(temp_downloads::table)
        .values(map)
        .execute(conn)
}

fn save_to_version_downloads(conn: &mut PgConnection) -> QueryResult<usize> {
    diesel::sql_query(
        r#"
            INSERT INTO version_downloads (version_id, date, downloads)
            SELECT versions.id, temp_downloads.date, temp_downloads.downloads FROM temp_downloads
                INNER JOIN crates ON crates.name = temp_downloads.name
                INNER JOIN versions ON versions.num = temp_downloads.version
                       AND versions.crate_id = crates.id
            ON CONFLICT (version_id, date) DO UPDATE SET downloads = version_downloads.downloads + EXCLUDED.downloads;
        "#,
    )
    .execute(conn)
}