Implement migratable database stores

Add migratable key listing for SQLite and PostgreSQL stores so benchmark
fixtures can copy persisted node state between database backends.

AI-assisted-by: OpenAI Codex

Author: benthecarman

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;
@@ -339,6 +341,19 @@ 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 {
+			run_on_internal_runtime(runtime, async move { inner.list_all_keys_internal().await })
+				.await
+		}
+	}
+}
+
 struct PostgresStoreInner {
 	pool: SmallPool,
 	config: Config,
@@ -713,6 +728,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> {
@@ -893,6 +927,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,35 @@ 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 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::new();
+		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 +725,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();