From b6538221bb0be3a0c106fcc26a434262b0cf969e Mon Sep 17 00:00:00 2001 From: Justin Chang Date: Wed, 2 Aug 2023 18:54:40 -0700 Subject: [PATCH 1/2] Migrations in code --- Cargo.lock | 1 + .../nft-metadata-crawler-parser/Cargo.toml | 1 + .../nft-metadata-crawler-parser/src/lib.rs | 12 ---------- .../src/utils/database.rs | 22 ++++++++++++++++++- .../nft-metadata-crawler-parser/src/worker.rs | 9 ++++++-- 5 files changed, 30 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9ffbcb829a5f5..542c7b5a69054 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9564,6 +9564,7 @@ dependencies = [ "crossbeam-channel", "csv", "diesel", + "diesel_migrations", "field_count", "futures", "google-cloud-pubsub", diff --git a/ecosystem/nft-metadata-crawler-parser/Cargo.toml b/ecosystem/nft-metadata-crawler-parser/Cargo.toml index a741435cce681..b0c1422b4cdb1 100644 --- a/ecosystem/nft-metadata-crawler-parser/Cargo.toml +++ b/ecosystem/nft-metadata-crawler-parser/Cargo.toml @@ -23,6 +23,7 @@ clap = { workspace = true } crossbeam-channel = { workspace = true } csv = { workspace = true } diesel = { workspace = true } +diesel_migrations = { workspace = true } field_count = { workspace = true } futures = { workspace = true } google-cloud-pubsub = { workspace = true } diff --git a/ecosystem/nft-metadata-crawler-parser/src/lib.rs b/ecosystem/nft-metadata-crawler-parser/src/lib.rs index 1a0c3515ca408..daf2b15041f8e 100644 --- a/ecosystem/nft-metadata-crawler-parser/src/lib.rs +++ b/ecosystem/nft-metadata-crawler-parser/src/lib.rs @@ -5,20 +5,8 @@ pub mod schema; pub mod utils; pub mod worker; -use diesel::{ - pg::PgConnection, - r2d2::{ConnectionManager, Pool}, -}; use reqwest::{header, Client}; -/// Establishes a connection pool to Postgres -pub fn establish_connection_pool(database_url: String) -> Pool> { - let manager = ConnectionManager::::new(database_url); - Pool::builder() - .build(manager) - .expect("Failed to create pool.") -} - /// HEAD request to get MIME type and size of content pub async fn get_uri_metadata(url: String) -> anyhow::Result<(String, u32)> { let client = Client::new(); diff --git a/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs b/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs index 31f57d190507d..49d0fd667938d 100644 --- a/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs +++ b/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs @@ -3,12 +3,32 @@ use crate::{models::nft_metadata_crawler_uris::NFTMetadataCrawlerURIs, schema}; use anyhow::Context; use diesel::{ - r2d2::{ConnectionManager, PooledConnection}, + r2d2::{ConnectionManager, Pool, PooledConnection}, upsert::excluded, ExpressionMethods, PgConnection, RunQueryDsl, }; +use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; use tracing::debug; +pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!(); + +/// Establishes a connection pool to Postgres +pub fn establish_connection_pool(database_url: String) -> Pool> { + let manager = ConnectionManager::::new(database_url); + let pool = Pool::builder() + .build(manager) + .expect("Failed to create pool."); + pool +} + +/// Runs database migrations +pub fn run_migrations(pool: &Pool>) { + pool.get() + .expect("[NFT Metadata Crawler] Could not get connection for migrations") + .run_pending_migrations(MIGRATIONS) + .expect("[NFT Metadata Crawler] migrations failed!"); +} + /// Upserts URIs into database pub fn upsert_uris( conn: &mut PooledConnection>, diff --git a/ecosystem/nft-metadata-crawler-parser/src/worker.rs b/ecosystem/nft-metadata-crawler-parser/src/worker.rs index 5a23dfcc52200..54e9c758b0141 100644 --- a/ecosystem/nft-metadata-crawler-parser/src/worker.rs +++ b/ecosystem/nft-metadata-crawler-parser/src/worker.rs @@ -1,13 +1,12 @@ // Copyright © Aptos Foundation use crate::{ - establish_connection_pool, models::{ nft_metadata_crawler_uris::NFTMetadataCrawlerURIs, nft_metadata_crawler_uris_query::NFTMetadataCrawlerURIsQuery, }, utils::{ - database::upsert_uris, + database::{establish_connection_pool, run_migrations, upsert_uris}, gcs::{write_image_to_gcs, write_json_to_gcs}, image_optimizer::ImageOptimizer, json_parser::JSONParser, @@ -112,7 +111,13 @@ async fn spawn_parser( impl RunnableConfig for ParserConfig { /// Main driver function that establishes a connection to Pubsub and parses the Pubsub entries in parallel async fn run(&self) -> anyhow::Result<()> { + info!("[NFT Metadata Crawler] Connecting to database"); let pool = establish_connection_pool(self.database_url.clone()); + info!("[NFT Metadata Crawler] Database connection successful"); + + info!("[NFT Metadata Crawler] Running migrations"); + run_migrations(&pool); + info!("[NFT Metadata Crawler] Finished migrations"); std::env::set_var( "GOOGLE_APPLICATION_CREDENTIALS", From a2c1208b021f7bdd9eabde9b07a532b22c4a96f5 Mon Sep 17 00:00:00 2001 From: Justin Chang Date: Wed, 2 Aug 2023 21:20:23 -0700 Subject: [PATCH 2/2] lint --- ecosystem/nft-metadata-crawler-parser/src/utils/database.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs b/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs index 49d0fd667938d..324db90e88e73 100644 --- a/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs +++ b/ecosystem/nft-metadata-crawler-parser/src/utils/database.rs @@ -15,10 +15,9 @@ pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!(); /// Establishes a connection pool to Postgres pub fn establish_connection_pool(database_url: String) -> Pool> { let manager = ConnectionManager::::new(database_url); - let pool = Pool::builder() + Pool::builder() .build(manager) - .expect("Failed to create pool."); - pool + .expect("Failed to create pool.") } /// Runs database migrations