[NFT Metadata Crawler] DB migrations in code (#9451)

aptos-labs · Aug 3, 2023 · 8fe69b2 · 8fe69b2
1 parent efb8348
commit 8fe69b2
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 15 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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
@@ -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<ConnectionManager<PgConnection>> {
-    let manager = ConnectionManager::<PgConnection>::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
@@ -3,12 +3,31 @@
 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<ConnectionManager<PgConnection>> {
+    let manager = ConnectionManager::<PgConnection>::new(database_url);
+    Pool::builder()
+        .build(manager)
+        .expect("Failed to create pool.")
+}
+
+/// Runs database migrations
+pub fn run_migrations(pool: &Pool<ConnectionManager<PgConnection>>) {
+    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<ConnectionManager<PgConnection>>,

diff --git 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",