Skip to content

Instantly share code, notes, and snippets.

@Turbo87
Created January 26, 2024 17:41
Show Gist options
  • Save Turbo87/a8465a7a62189d2ce9f99d547fee460a to your computer and use it in GitHub Desktop.
Save Turbo87/a8465a7a62189d2ce9f99d547fee460a to your computer and use it in GitHub Desktop.
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)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment