Refactor catalog manager to be fully async (#34)

* refactor catalog manager to be async

* remove unnecessary multi_thread flavor

* move `block_on` to datafusion module

Author: zfarrell

src/bin/server.rs

@@ -34,7 +34,7 @@ async fn main() -> Result<()> {
     tracing::info!("Configuration '{}' loaded successfully", &cli.config);
 
     // Initialize engine from config
-    let engine = HotDataEngine::from_config(&config)?;
+    let engine = HotDataEngine::from_config(&config).await?;
 
     tracing::info!("Engine initialized");
 
@@ -55,7 +55,7 @@ async fn main() -> Result<()> {
     server.await?;
 
     // Explicitly shutdown engine to close catalog connection
-    if let Err(e) = engine.shutdown() {
+    if let Err(e) = engine.shutdown().await {
         tracing::error!("Error during engine shutdown: {}", e);
     }
 

src/catalog/manager.rs

@@ -1,8 +1,8 @@
 use anyhow::Result;
+use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use sqlx::FromRow;
 use std::fmt::Debug;
-use tokio::task::block_in_place;
 
 #[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
 pub struct ConnectionInfo {
@@ -24,61 +24,54 @@ pub struct TableInfo {
     pub arrow_schema_json: Option<String>,
 }
 
-/// Blocking helper for async operations.
-///
-/// Uses `block_in_place` to avoid blocking the tokio runtime when calling
-/// async code from a sync context.
-pub fn block_on<F, T>(f: F) -> Result<T>
-where
-    F: std::future::Future<Output = Result<T>>,
-{
-    block_in_place(|| tokio::runtime::Handle::current().block_on(f))
-}
-
-/// Synchronous interface for catalog operations.
-///
-/// This trait uses blocking methods for compatibility with `dyn` trait objects.
-/// Implementations should use [`block_on`] to wrap async database operations.
+/// Async interface for catalog operations.
+#[async_trait]
 pub trait CatalogManager: Debug + Send + Sync {
     /// Close the catalog connection. This is idempotent and can be called multiple times.
-    fn close(&self) -> Result<()> {
+    async fn close(&self) -> Result<()> {
         // Default implementation does nothing - sqlx pools handle cleanup automatically
         Ok(())
     }
 
     /// Apply any pending schema migrations. Should be idempotent.
-    fn run_migrations(&self) -> Result<()>;
+    async fn run_migrations(&self) -> Result<()>;
 
-    fn list_connections(&self) -> Result<Vec<ConnectionInfo>>;
-    fn add_connection(&self, name: &str, source_type: &str, config_json: &str) -> Result<i32>;
-    fn get_connection(&self, name: &str) -> Result<Option<ConnectionInfo>>;
-    fn add_table(
+    async fn list_connections(&self) -> Result<Vec<ConnectionInfo>>;
+    async fn add_connection(&self, name: &str, source_type: &str, config_json: &str)
+        -> Result<i32>;
+    async fn get_connection(&self, name: &str) -> Result<Option<ConnectionInfo>>;
+    async fn add_table(
         &self,
         connection_id: i32,
         schema_name: &str,
         table_name: &str,
         arrow_schema_json: &str,
     ) -> Result<i32>;
-    fn list_tables(&self, connection_id: Option<i32>) -> Result<Vec<TableInfo>>;
-    fn get_table(
+    async fn list_tables(&self, connection_id: Option<i32>) -> Result<Vec<TableInfo>>;
+    async fn get_table(
         &self,
         connection_id: i32,
         schema_name: &str,
         table_name: &str,
     ) -> Result<Option<TableInfo>>;
-    fn update_table_sync(&self, table_id: i32, parquet_path: &str, state_path: &str) -> Result<()>;
+    async fn update_table_sync(
+        &self,
+        table_id: i32,
+        parquet_path: &str,
+        state_path: &str,
+    ) -> Result<()>;
 
     /// Clear table cache metadata (set paths to NULL) without deleting files.
-    fn clear_table_cache_metadata(
+    async fn clear_table_cache_metadata(
         &self,
         connection_id: i32,
         schema_name: &str,
         table_name: &str,
     ) -> Result<TableInfo>;
 
     /// Clear cache metadata for all tables in a connection (set paths to NULL).
-    fn clear_connection_cache_metadata(&self, name: &str) -> Result<()>;
+    async fn clear_connection_cache_metadata(&self, name: &str) -> Result<()>;
 
     /// Delete connection and all associated table rows from metadata.
-    fn delete_connection(&self, name: &str) -> Result<()>;
+    async fn delete_connection(&self, name: &str) -> Result<()>;
 }

src/catalog/postgres_manager.rs

@@ -1,7 +1,8 @@
 use crate::catalog::backend::CatalogBackend;
-use crate::catalog::manager::{block_on, CatalogManager, ConnectionInfo, TableInfo};
+use crate::catalog::manager::{CatalogManager, ConnectionInfo, TableInfo};
 use crate::catalog::migrations::{run_migrations, CatalogMigrations};
 use anyhow::Result;
+use async_trait::async_trait;
 use sqlx::postgres::PgPoolOptions;
 use sqlx::{PgPool, Postgres};
 use std::fmt::{self, Debug, Formatter};
@@ -11,11 +12,7 @@ pub struct PostgresCatalogManager {
 }
 
 impl PostgresCatalogManager {
-    pub fn new(connection_string: &str) -> Result<Self> {
-        block_on(Self::new_async(connection_string))
-    }
-
-    async fn new_async(connection_string: &str) -> Result<Self> {
+    pub async fn new(connection_string: &str) -> Result<Self> {
         let pool = PgPoolOptions::new()
             .max_connections(5)
             .connect(connection_string)
@@ -97,79 +94,86 @@ impl CatalogMigrations for PostgresMigrationBackend {
     }
 }
 
+#[async_trait]
 impl CatalogManager for PostgresCatalogManager {
-    fn run_migrations(&self) -> Result<()> {
-        block_on(run_migrations::<PostgresMigrationBackend>(
-            self.backend.pool(),
-        ))
+    async fn run_migrations(&self) -> Result<()> {
+        run_migrations::<PostgresMigrationBackend>(self.backend.pool()).await
     }
 
-    fn list_connections(&self) -> Result<Vec<ConnectionInfo>> {
-        block_on(self.backend.list_connections())
+    async fn list_connections(&self) -> Result<Vec<ConnectionInfo>> {
+        self.backend.list_connections().await
     }
 
-    fn add_connection(&self, name: &str, source_type: &str, config_json: &str) -> Result<i32> {
-        block_on(self.backend.add_connection(name, source_type, config_json))
+    async fn add_connection(
+        &self,
+        name: &str,
+        source_type: &str,
+        config_json: &str,
+    ) -> Result<i32> {
+        self.backend
+            .add_connection(name, source_type, config_json)
+            .await
     }
 
-    fn get_connection(&self, name: &str) -> Result<Option<ConnectionInfo>> {
-        block_on(self.backend.get_connection(name))
+    async fn get_connection(&self, name: &str) -> Result<Option<ConnectionInfo>> {
+        self.backend.get_connection(name).await
     }
 
-    fn add_table(
+    async fn add_table(
         &self,
         connection_id: i32,
         schema_name: &str,
         table_name: &str,
         arrow_schema_json: &str,
     ) -> Result<i32> {
-        block_on(
-            self.backend
-                .add_table(connection_id, schema_name, table_name, arrow_schema_json),
-        )
+        self.backend
+            .add_table(connection_id, schema_name, table_name, arrow_schema_json)
+            .await
     }
 
-    fn list_tables(&self, connection_id: Option<i32>) -> Result<Vec<TableInfo>> {
-        block_on(self.backend.list_tables(connection_id))
+    async fn list_tables(&self, connection_id: Option<i32>) -> Result<Vec<TableInfo>> {
+        self.backend.list_tables(connection_id).await
     }
 
-    fn get_table(
+    async fn get_table(
         &self,
         connection_id: i32,
         schema_name: &str,
         table_name: &str,
     ) -> Result<Option<TableInfo>> {
-        block_on(
-            self.backend
-                .get_table(connection_id, schema_name, table_name),
-        )
+        self.backend
+            .get_table(connection_id, schema_name, table_name)
+            .await
     }
 
-    fn update_table_sync(&self, table_id: i32, parquet_path: &str, state_path: &str) -> Result<()> {
-        block_on(
-            self.backend
-                .update_table_sync(table_id, parquet_path, state_path),
-        )
+    async fn update_table_sync(
+        &self,
+        table_id: i32,
+        parquet_path: &str,
+        state_path: &str,
+    ) -> Result<()> {
+        self.backend
+            .update_table_sync(table_id, parquet_path, state_path)
+            .await
     }
 
-    fn clear_table_cache_metadata(
+    async fn clear_table_cache_metadata(
         &self,
         connection_id: i32,
         schema_name: &str,
         table_name: &str,
     ) -> Result<TableInfo> {
-        block_on(
-            self.backend
-                .clear_table_cache_metadata(connection_id, schema_name, table_name),
-        )
+        self.backend
+            .clear_table_cache_metadata(connection_id, schema_name, table_name)
+            .await
     }
 
-    fn clear_connection_cache_metadata(&self, name: &str) -> Result<()> {
-        block_on(self.backend.clear_connection_cache_metadata(name))
+    async fn clear_connection_cache_metadata(&self, name: &str) -> Result<()> {
+        self.backend.clear_connection_cache_metadata(name).await
     }
 
-    fn delete_connection(&self, name: &str) -> Result<()> {
-        block_on(self.backend.delete_connection(name))
+    async fn delete_connection(&self, name: &str) -> Result<()> {
+        self.backend.delete_connection(name).await
     }
 }