Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[NFT Metadata Crawler] DB migrations in code #9451

Merged
merged 2 commits into from
Aug 3, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions ecosystem/nft-metadata-crawler-parser/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
12 changes: 0 additions & 12 deletions ecosystem/nft-metadata-crawler-parser/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
21 changes: 20 additions & 1 deletion ecosystem/nft-metadata-crawler-parser/src/utils/database.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>>,
Expand Down
9 changes: 7 additions & 2 deletions ecosystem/nft-metadata-crawler-parser/src/worker.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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",
Expand Down