Merge pull request #945 from benthecarman/migrate-dbs

Implement `MigratableKVStore` for all KV Stores

Author: tnull

Cargo.toml

@@ -92,6 +92,7 @@ winapi = { version = "0.3", features = ["winbase"] }
 
 [dev-dependencies]
 lightning = { git = "https://github.com/lightningdevkit/rust-lightning", rev = "3dfcc4cca1866c5e5d4d4eaf3b82e09584e2ce5c", features = ["std", "_test_utils"] }
+lightning-persister = { git = "https://github.com/lightningdevkit/rust-lightning", rev = "3dfcc4cca1866c5e5d4d4eaf3b82e09584e2ce5c", features = ["tokio"] }
 rand = { version = "0.9.2", default-features = false, features = ["std", "thread_rng", "os_rng"] }
 proptest = "1.0.0"
 regex = "1.5.6"

src/io/in_memory_store.rs

@@ -11,7 +11,9 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Mutex;
 
 use lightning::io;
-use lightning::util::persist::{KVStore, PageToken, PaginatedKVStore, PaginatedListResponse};
+use lightning::util::persist::{
+	KVStore, MigratableKVStore, PageToken, PaginatedKVStore, PaginatedListResponse,
+};
 
 const IN_MEMORY_PAGE_SIZE: usize = 50;
 
@@ -96,6 +98,28 @@ impl InMemoryStore {
 			hash_map::Entry::Vacant(_) => Ok(Vec::new()),
 		}
 	}
+
+	fn list_all_keys_internal(&self) -> io::Result<Vec<(String, String, String)>> {
+		let persisted_lock = self.persisted_bytes.lock().unwrap();
+		let capacity = persisted_lock.values().map(|entries| entries.len()).sum();
+		let mut keys = Vec::with_capacity(capacity);
+
+		for (prefixed_namespace, namespace_entries) in persisted_lock.iter() {
+			let (primary_namespace, secondary_namespace) =
+				prefixed_namespace.split_once('/').ok_or_else(|| {
+					io::Error::new(io::ErrorKind::InvalidData, "Invalid namespace format")
+				})?;
+			for key in namespace_entries.keys() {
+				keys.push((
+					primary_namespace.to_string(),
+					secondary_namespace.to_string(),
+					key.clone(),
+				));
+			}
+		}
+
+		Ok(keys)
+	}
 }
 
 impl KVStore for InMemoryStore {
@@ -187,5 +211,40 @@ impl PaginatedKVStore for InMemoryStore {
 	}
 }
 
+impl MigratableKVStore for InMemoryStore {
+	fn list_all_keys(
+		&self,
+	) -> impl Future<Output = Result<Vec<(String, String, String)>, io::Error>> + 'static + Send {
+		let res = self.list_all_keys_internal();
+		async move { res }
+	}
+}
+
 unsafe impl Sync for InMemoryStore {}
 unsafe impl Send for InMemoryStore {}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[tokio::test]
+	async fn in_memory_store_list_all_keys() {
+		let store = InMemoryStore::new();
+
+		KVStore::write(&store, "ns_a", "sub_a", "key_a", vec![1u8]).await.unwrap();
+		KVStore::write(&store, "ns_a", "sub_b", "key_b", vec![2u8]).await.unwrap();
+		KVStore::write(&store, "ns_b", "", "key_c", vec![3u8]).await.unwrap();
+
+		let mut keys = MigratableKVStore::list_all_keys(&store).await.unwrap();
+		keys.sort();
+
+		assert_eq!(
+			keys,
+			vec![
+				("ns_a".to_string(), "sub_a".to_string(), "key_a".to_string()),
+				("ns_a".to_string(), "sub_b".to_string(), "key_b".to_string()),
+				("ns_b".to_string(), "".to_string(), "key_c".to_string()),
+			]
+		);
+	}
+}

src/io/postgres_store/mod.rs

@@ -12,7 +12,9 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 
 use lightning::io;
-use lightning::util::persist::{KVStore, PageToken, PaginatedKVStore, PaginatedListResponse};
+use lightning::util::persist::{
+	KVStore, MigratableKVStore, PageToken, PaginatedKVStore, PaginatedListResponse,
+};
 use lightning_types::string::PrintableString;
 use native_tls::TlsConnector;
 use postgres_native_tls::MakeTlsConnector;
@@ -351,6 +353,24 @@ impl PaginatedKVStore for PostgresStore {
 	}
 }
 
+impl MigratableKVStore for PostgresStore {
+	fn list_all_keys(
+		&self,
+	) -> impl Future<Output = Result<Vec<(String, String, String)>, io::Error>> + 'static + Send {
+		let inner = Arc::clone(&self.inner);
+		let runtime = self.internal_runtime();
+		async move {
+			let task = runtime.spawn(async move { inner.list_all_keys_internal().await });
+			task.await.map_err(|e| {
+				io::Error::new(
+					io::ErrorKind::Other,
+					format!("PostgreSQL runtime task failed: {}", e),
+				)
+			})?
+		}
+	}
+}
+
 struct PostgresStoreInner {
 	pool: SmallPool,
 	config: Config,
@@ -725,6 +745,25 @@ impl PostgresStoreInner {
 		Ok(keys)
 	}
 
+	async fn list_all_keys_internal(&self) -> io::Result<Vec<(String, String, String)>> {
+		let sql = format!(
+			"SELECT primary_namespace, secondary_namespace, key FROM {}",
+			self.kv_table_name_sql
+		);
+
+		let err_map = |e: PgError| {
+			let msg = format!("Failed to retrieve queried rows: {e}");
+			io::Error::new(io::ErrorKind::Other, msg)
+		};
+
+		let mut locked = self.locked_client().await?;
+		let rows = query_with_retry!(self, locked, err_map, locked.query(sql.as_str(), &[]))?;
+
+		let keys: Vec<(String, String, String)> =
+			rows.iter().map(|row| (row.get(0), row.get(1), row.get(2))).collect();
+		Ok(keys)
+	}
+
 	async fn list_paginated_internal(
 		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
 	) -> io::Result<PaginatedListResponse> {
@@ -904,6 +943,29 @@ mod tests {
 		cleanup_store(&store_1).await;
 	}
 
+	#[tokio::test(flavor = "multi_thread")]
+	async fn test_postgres_store_list_all_keys() {
+		let store = create_test_store("test_pg_list_all_keys").await;
+
+		KVStore::write(&store, "ns_a", "sub_a", "key_a", vec![1u8]).await.unwrap();
+		KVStore::write(&store, "ns_a", "sub_b", "key_b", vec![2u8]).await.unwrap();
+		KVStore::write(&store, "ns_b", "", "key_c", vec![3u8]).await.unwrap();
+
+		let mut keys = MigratableKVStore::list_all_keys(&store).await.unwrap();
+		keys.sort();
+
+		assert_eq!(
+			keys,
+			vec![
+				("ns_a".to_string(), "sub_a".to_string(), "key_a".to_string()),
+				("ns_a".to_string(), "sub_b".to_string(), "key_b".to_string()),
+				("ns_b".to_string(), "".to_string(), "key_c".to_string()),
+			]
+		);
+
+		cleanup_store(&store).await;
+	}
+
 	async fn kill_connection(store: &PostgresStore) {
 		// Terminate every backend in the pool so the next op deterministically
 		// hits a closed connection regardless of which slot `get` selects.

src/io/sqlite_store/mod.rs

@@ -14,7 +14,9 @@ use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 
 use lightning::io;
-use lightning::util::persist::{KVStore, PageToken, PaginatedKVStore, PaginatedListResponse};
+use lightning::util::persist::{
+	KVStore, MigratableKVStore, PageToken, PaginatedKVStore, PaginatedListResponse,
+};
 use lightning_types::string::PrintableString;
 use rusqlite::{named_params, Connection};
 
@@ -202,6 +204,21 @@ impl PaginatedKVStore for SqliteStore {
 	}
 }
 
+impl MigratableKVStore for SqliteStore {
+	fn list_all_keys(
+		&self,
+	) -> impl Future<Output = Result<Vec<(String, String, String)>, io::Error>> + 'static + Send {
+		let inner = Arc::clone(&self.inner);
+		let fut = tokio::task::spawn_blocking(move || inner.list_all_keys_internal());
+		async move {
+			fut.await.unwrap_or_else(|e| {
+				let msg = format!("Failed to IO operation due join error: {}", e);
+				Err(io::Error::new(io::ErrorKind::Other, msg))
+			})
+		}
+	}
+}
+
 struct SqliteStoreInner {
 	connection: Arc<Mutex<Connection>>,
 	data_dir: PathBuf,
@@ -486,6 +503,42 @@ impl SqliteStoreInner {
 		Ok(keys)
 	}
 
+	fn list_all_keys_internal(&self) -> io::Result<Vec<(String, String, String)>> {
+		let locked_conn = self.connection.lock().expect("lock");
+
+		let sql = format!(
+			"SELECT primary_namespace, secondary_namespace, key FROM {}",
+			self.kv_table_name
+		);
+		let count_sql = format!("SELECT COUNT(*) FROM {}", self.kv_table_name);
+		let count: usize =
+			locked_conn.query_row(&count_sql, [], |row| row.get(0)).map_err(|e| {
+				let msg = format!("Failed to count rows: {}", e);
+				io::Error::new(io::ErrorKind::Other, msg)
+			})?;
+
+		let mut stmt = locked_conn.prepare_cached(&sql).map_err(|e| {
+			let msg = format!("Failed to prepare statement: {}", e);
+			io::Error::new(io::ErrorKind::Other, msg)
+		})?;
+
+		let mut keys = Vec::with_capacity(count);
+		let rows_iter =
+			stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?))).map_err(|e| {
+				let msg = format!("Failed to retrieve queried rows: {}", e);
+				io::Error::new(io::ErrorKind::Other, msg)
+			})?;
+
+		for key in rows_iter {
+			keys.push(key.map_err(|e| {
+				let msg = format!("Failed to retrieve queried rows: {}", e);
+				io::Error::new(io::ErrorKind::Other, msg)
+			})?);
+		}
+
+		Ok(keys)
+	}
+
 	fn list_paginated_internal(
 		&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
 	) -> io::Result<PaginatedListResponse> {
@@ -679,6 +732,34 @@ mod tests {
 		do_test_store(&store_0, &store_1)
 	}
 
+	#[tokio::test]
+	async fn test_sqlite_store_list_all_keys() {
+		let mut temp_path = random_storage_path();
+		temp_path.push("test_sqlite_store_list_all_keys");
+		let store = SqliteStore::new(
+			temp_path,
+			Some("test_db".to_string()),
+			Some("test_table".to_string()),
+		)
+		.unwrap();
+
+		KVStore::write(&store, "ns_a", "sub_a", "key_a", vec![1u8]).await.unwrap();
+		KVStore::write(&store, "ns_a", "sub_b", "key_b", vec![2u8]).await.unwrap();
+		KVStore::write(&store, "ns_b", "", "key_c", vec![3u8]).await.unwrap();
+
+		let mut keys = MigratableKVStore::list_all_keys(&store).await.unwrap();
+		keys.sort();
+
+		assert_eq!(
+			keys,
+			vec![
+				("ns_a".to_string(), "sub_a".to_string(), "key_a".to_string()),
+				("ns_a".to_string(), "sub_b".to_string(), "key_b".to_string()),
+				("ns_b".to_string(), "".to_string(), "key_c".to_string()),
+			]
+		);
+	}
+
 	#[tokio::test]
 	async fn test_sqlite_store_paginated_listing() {
 		let mut temp_path = random_storage_path();