Add migrations tests

Author: tankyleo

rust/impls/src/postgres_store.rs

@@ -28,12 +28,18 @@ const VALUE_COLUMN: &str = "value";
 const VERSION_COLUMN: &str = "version";
 
 const DB_VERSION_COLUMN: &str = "db_version";
+#[cfg(test)]
+const MIGRATION_LOG_COLUMN: &str = "upgrade_from";
 
 const CHECK_DB_STMT: &str = "SELECT 1 FROM pg_database WHERE datname = $1";
 const INIT_DB_CMD: &str = "CREATE DATABASE";
+#[cfg(test)]
+const DROP_DB_CMD: &str = "DROP DATABASE";
 const GET_VERSION_STMT: &str = "SELECT db_version FROM vss_db_version;";
 const UPDATE_VERSION_STMT: &str = "UPDATE vss_db_version SET db_version=$1;";
 const LOG_MIGRATION_STMT: &str = "INSERT INTO vss_db_upgrades VALUES($1);";
+#[cfg(test)]
+const GET_MIGRATION_LOG_STMT: &str = "SELECT upgrade_from FROM vss_db_upgrades;";
 
 const MIGRATIONS: &[&str] = &[
 	"CREATE TABLE vss_db_version (db_version INTEGER);",
@@ -52,6 +58,8 @@ const MIGRATIONS: &[&str] = &[
 	    PRIMARY KEY (user_token, store_id, key)
 	);",
 ];
+#[cfg(test)]
+const DUMMY_MIGRATION: &str = "SELECT 1 WHERE FALSE;";
 
 /// The maximum number of key versions that can be returned in a single page.
 ///
@@ -102,6 +110,31 @@ async fn initialize_vss_database(postgres_endpoint: &str, db_name: &str) -> Resu
 	Ok(())
 }
 
+#[cfg(test)]
+async fn drop_database(postgres_endpoint: &str, db_name: &str) -> Result<(), Error> {
+	let postgres_dsn = format!("{}/{}", postgres_endpoint, "postgres");
+	let (client, connection) = tokio_postgres::connect(&postgres_dsn, NoTls)
+		.await
+		.map_err(|e| Error::new(ErrorKind::Other, format!("Connection error: {}", e)))?;
+	// Connection must be driven on a separate task, and will resolve when the client is dropped
+	tokio::spawn(async move {
+		if let Err(e) = connection.await {
+			eprintln!("Connection error: {}", e);
+		}
+	});
+
+	let drop_database_statement = format!("{} {};", DROP_DB_CMD, db_name);
+	let num_rows = client.execute(&drop_database_statement, &[]).await.map_err(|e| {
+		Error::new(
+			ErrorKind::Other,
+			format!("Failed to drop database {}: {}", db_name, e),
+		)
+	})?;
+	assert_eq!(num_rows, 0);
+
+	Ok(())
+}
+
 impl PostgresBackendImpl {
 	/// Constructs a [`PostgresBackendImpl`] using `dsn` for PostgreSQL connection information.
 	pub async fn new(postgres_endpoint: &str, db_name: &str) -> Result<Self, Error> {
@@ -125,12 +158,13 @@ impl PostgresBackendImpl {
 			.map_err(|e| Error::new(ErrorKind::Other, format!("Failed to build Pool: {}", e)))?;
 		let postgres_backend = PostgresBackendImpl { pool };
 
-		postgres_backend.migrate_vss_database().await?;
+		#[cfg(not(test))]
+		postgres_backend.migrate_vss_database(MIGRATIONS).await?;
 
 		Ok(postgres_backend)
 	}
 
-	async fn migrate_vss_database(&self) -> Result<(), Error> {
+	async fn migrate_vss_database(&self, migrations: &[&str]) -> Result<(usize, usize), Error> {
 		let mut conn = self.pool.get().await.map_err(|e| {
 			Error::new(
 				ErrorKind::Other,
@@ -162,16 +196,16 @@ impl PostgresBackendImpl {
 			.await
 			.map_err(|e| Error::new(ErrorKind::Other, format!("Transaction start error: {}", e)))?;
 
-		if migration_start == MIGRATIONS.len() {
+		if migration_start == migrations.len() {
 			// No migrations needed, we are done
-			return Ok(());
-		} else if migration_start > MIGRATIONS.len() {
+			return Ok((migration_start, migrations.len()));
+		} else if migration_start > migrations.len() {
 			panic!("We do not allow downgrades");
 		}
 
-		println!("Applying migration(s) {} through {}", migration_start, MIGRATIONS.len() - 1);
+		println!("Applying migration(s) {} through {}", migration_start, migrations.len() - 1);
 
-		for (idx, &stmt) in (&MIGRATIONS[migration_start..]).iter().enumerate() {
+		for (idx, &stmt) in (&migrations[migration_start..]).iter().enumerate() {
 			let _num_rows = tx.execute(stmt, &[]).await.map_err(|e| {
 				Error::new(
 					ErrorKind::Other,
@@ -197,7 +231,7 @@ impl PostgresBackendImpl {
 		assert_eq!(num_rows, 1, "LOG_MIGRATION_STMT should only add one row at a time");
 
 		let next_migration_start =
-			i32::try_from(MIGRATIONS.len()).expect("Length is definitely smaller than i32::MAX");
+			i32::try_from(migrations.len()).expect("Length is definitely smaller than i32::MAX");
 		let num_rows =
 			tx.execute(UPDATE_VERSION_STMT, &[&next_migration_start]).await.map_err(|e| {
 				Error::new(
@@ -214,7 +248,21 @@ impl PostgresBackendImpl {
 			Error::new(ErrorKind::Other, format!("Transaction commit error: {}", e))
 		})?;
 
-		Ok(())
+		Ok((migration_start, migrations.len()))
+	}
+
+	#[cfg(test)]
+	async fn get_schema_version(&self) -> usize {
+		let conn = self.pool.get().await.unwrap();
+		let row = conn.query_one(GET_VERSION_STMT, &[]).await.unwrap();
+		usize::try_from(row.get::<&str, i32>(DB_VERSION_COLUMN)).unwrap()
+	}
+
+	#[cfg(test)]
+	async fn get_upgrades_list(&self) -> Vec<usize> {
+		let conn = self.pool.get().await.unwrap();
+		let rows = conn.query(GET_MIGRATION_LOG_STMT, &[]).await.unwrap();
+		rows.iter().map(|row| usize::try_from(row.get::<&str, i32>(MIGRATION_LOG_COLUMN)).unwrap()).collect()
 	}
 
 	fn build_vss_record(&self, user_token: String, store_id: String, kv: KeyValue) -> VssDbRecord {
@@ -568,23 +616,104 @@ mod tests {
 	use crate::postgres_store::PostgresBackendImpl;
 	use api::define_kv_store_tests;
 	use tokio::sync::OnceCell;
+	use super::{MIGRATIONS, DUMMY_MIGRATION, drop_database};
+
+	const POSTGRES_ENDPOINT: &str = "postgresql://postgres:postgres@localhost:5432";
+	const MIGRATIONS_START: usize = 0;
+	const MIGRATIONS_END: usize = MIGRATIONS.len();
 
 	static START: OnceCell<()> = OnceCell::const_new();
 
 	define_kv_store_tests!(PostgresKvStoreTest, PostgresBackendImpl, {
+		let db_name = "postgres_kv_store_tests";
 		START
 			.get_or_init(|| async {
-				// Initialize the database once, and have other threads wait
-				PostgresBackendImpl::new(
-					"postgresql://postgres:postgres@localhost:5432",
-					"postgres",
-				)
-				.await
-				.unwrap();
+				let _ = drop_database(POSTGRES_ENDPOINT, db_name).await;
+				let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+				let (start, end) = store.migrate_vss_database(MIGRATIONS).await.unwrap();
+				assert_eq!(start, MIGRATIONS_START);
+				assert_eq!(end, MIGRATIONS_END);
 			})
 			.await;
-		PostgresBackendImpl::new("postgresql://postgres:postgres@localhost:5432", "postgres")
-			.await
-			.unwrap()
+		let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+		let (start, end) = store.migrate_vss_database(MIGRATIONS).await.unwrap();
+		assert_eq!(start, MIGRATIONS_END);
+		assert_eq!(end, MIGRATIONS_END);
+		assert_eq!(store.get_upgrades_list().await, [MIGRATIONS_START]);
+		assert_eq!(store.get_schema_version().await, MIGRATIONS_END);
+		store
 	});
+
+	#[tokio::test]
+	#[should_panic(expected = "We do not allow downgrades")]
+	async fn panic_on_downgrade() {
+		let db_name = "panic_on_downgrade_test";
+		let _ = drop_database(POSTGRES_ENDPOINT, db_name).await;
+		{
+			let mut migrations = MIGRATIONS.to_vec();
+			migrations.push(DUMMY_MIGRATION);
+			let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+			let (start, end) = store.migrate_vss_database(&migrations).await.unwrap();
+			assert_eq!(start, MIGRATIONS_START);
+			assert_eq!(end, MIGRATIONS_END + 1);
+		};
+		{
+			let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+			let _ = store.migrate_vss_database(MIGRATIONS).await.unwrap();
+		};
+	}
+
+	#[tokio::test]
+	async fn new_migrations_increments_upgrades() {
+		let db_name = "new_migrations_increments_upgrades_test";
+		let _ = drop_database(POSTGRES_ENDPOINT, db_name).await;
+		{
+			let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+			let (start, end) = store.migrate_vss_database(MIGRATIONS).await.unwrap();
+			assert_eq!(start, MIGRATIONS_START);
+			assert_eq!(end, MIGRATIONS_END);
+			assert_eq!(store.get_upgrades_list().await, [MIGRATIONS_START]);
+			assert_eq!(store.get_schema_version().await, MIGRATIONS_END);
+		};
+		{
+			let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+			let (start, end) = store.migrate_vss_database(MIGRATIONS).await.unwrap();
+			assert_eq!(start, MIGRATIONS_END);
+			assert_eq!(end, MIGRATIONS_END);
+			assert_eq!(store.get_upgrades_list().await, [MIGRATIONS_START]);
+			assert_eq!(store.get_schema_version().await, MIGRATIONS_END);
+		};
+
+		let mut migrations = MIGRATIONS.to_vec();
+		migrations.push(DUMMY_MIGRATION);
+		{
+			let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+			let (start, end) = store.migrate_vss_database(&migrations).await.unwrap();
+			assert_eq!(start, MIGRATIONS_END);
+			assert_eq!(end, MIGRATIONS_END + 1);
+			assert_eq!(store.get_upgrades_list().await, [MIGRATIONS_START, MIGRATIONS_END]);
+			assert_eq!(store.get_schema_version().await, MIGRATIONS_END + 1);
+		};
+
+		migrations.push(DUMMY_MIGRATION);
+		migrations.push(DUMMY_MIGRATION);
+		{
+			let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+			let (start, end) = store.migrate_vss_database(&migrations).await.unwrap();
+			assert_eq!(start, MIGRATIONS_END + 1);
+			assert_eq!(end, MIGRATIONS_END + 3);
+			assert_eq!(store.get_upgrades_list().await, [MIGRATIONS_START, MIGRATIONS_END, MIGRATIONS_END + 1]);
+			assert_eq!(store.get_schema_version().await, MIGRATIONS_END + 3);
+		};
+
+		{
+			let store = PostgresBackendImpl::new(POSTGRES_ENDPOINT, db_name).await.unwrap();
+			let list = store.get_upgrades_list().await;
+			assert_eq!(list, [MIGRATIONS_START, MIGRATIONS_END, MIGRATIONS_END + 1]);
+			let version = store.get_schema_version().await;
+			assert_eq!(version, MIGRATIONS_END + 3);
+		}
+
+		drop_database(POSTGRES_ENDPOINT, db_name).await.unwrap();
+	}
 }