feat: summary loading

Author: sdd

tanic-svc/Cargo.toml

@@ -33,3 +33,4 @@ names = "0.14.0"
 tokio-stream = { version = "0.1.17", features = ["sync"] }
 parquet = "54.0.0"
 indexmap = "2.7.1"
+futures = "0.3.31"

tanic-svc/src/iceberg_context.rs

@@ -1,23 +1,29 @@
 //! Iceberg Context
 
-use iceberg::{Catalog, NamespaceIdent};
-use iceberg_catalog_rest::{RestCatalog, RestCatalogConfig};
 use std::sync::Arc;
+use std::sync::RwLock;
+
+use futures::stream::StreamExt;
+use iceberg::{Catalog, NamespaceIdent, TableIdent};
+use iceberg_catalog_rest::{RestCatalog, RestCatalogConfig};
 use tokio::sync::mpsc::UnboundedSender;
 use tokio::sync::watch::Receiver;
-use std::sync::RwLock;
 use tokio::task::JoinHandle;
-use tokio_stream::{wrappers::WatchStream, StreamExt};
+use tokio_stream::wrappers::WatchStream;
 
 use tanic_core::config::ConnectionDetails;
 use tanic_core::message::{NamespaceDeets, TableDeets};
 use tanic_core::{Result, TanicError};
+use tokio::sync::mpsc::{channel, Receiver as MpscReceiver, Sender as MpscSender};
+use tokio_stream::wrappers::ReceiverStream;
 
 use crate::state::{TanicAction, TanicAppState, TanicIcebergState};
 
 type ActionTx = UnboundedSender<TanicAction>;
 type IceCtxRef = Arc<RwLock<IcebergContext>>;
 
+const JOB_STREAM_CONCURRENCY: usize = 1;
+
 #[derive(Debug, Default)]
 struct IcebergContext {
     connection_details: Option<ConnectionDetails>,
@@ -40,6 +46,13 @@ pub struct IcebergContextManager {
     state_ref: Arc<RwLock<TanicAppState>>,
 }
 
+#[derive(Debug)]
+enum IcebergTask {
+    Namespaces,
+    TablesForNamespace(NamespaceDeets),
+    SummaryForTable(TableDeets),
+}
+
 impl IcebergContextManager {
     pub fn new(action_tx: ActionTx, state_ref: Arc<RwLock<TanicAppState>>) -> Self {
         Self {
@@ -52,8 +65,16 @@ impl IcebergContextManager {
     pub async fn event_loop(&self, state_rx: Receiver<()>) -> Result<()> {
         let mut state_stream = WatchStream::new(state_rx);
 
-        while state_stream.next().await.is_some() {
+        let (job_queue_tx, job_queue_rx) = channel(10);
+
+        tokio::spawn({
+            let action_tx = self.action_tx.clone();
+            let job_queue_tx = job_queue_tx.clone();
+            let iceberg_ctx = self.iceberg_context.clone();
+            async move { Self::job_handler(job_queue_rx, job_queue_tx, action_tx, iceberg_ctx).await }
+        });
 
+        while state_stream.next().await.is_some() {
             let new_conn_details = {
                 let state = self.state_ref.read().unwrap();
 
@@ -64,36 +85,27 @@ impl IcebergContextManager {
                     TanicIcebergState::Exiting => {
                         break;
                     }
-                    _ => None
+                    _ => None,
                 }
             };
 
             if let Some(new_conn_details) = new_conn_details {
-                self.connect_to(&new_conn_details).await?;
-
-                let namespaces = {
-                    self.iceberg_context.read().unwrap().namespaces.clone()
-                };
+                self.connect_to(&new_conn_details, job_queue_tx.clone())
+                    .await?;
 
                 // begin crawl
-                for namespace in namespaces {
-
-                    // spawn a task to start populating the namespaces
-                    let action_tx = self.action_tx.clone();
-                    let ctx = self.iceberg_context.clone();
-
-                    // TODO: handle handle, lol
-                    let _jh = tokio::spawn(async move {
-                        Self::populate_tables(ctx, action_tx, namespace).await
-                    });
-                }
+                let _ = job_queue_tx.send(IcebergTask::Namespaces).await;
             }
         }
 
         Ok(())
     }
 
-    async fn connect_to(&self, new_conn_details: &ConnectionDetails) -> Result<()> {
+    async fn connect_to(
+        &self,
+        new_conn_details: &ConnectionDetails,
+        _job_queue_tx: MpscSender<IcebergTask>,
+    ) -> Result<()> {
         {
             let ctx = self.iceberg_context.read().unwrap();
             if let Some(ref existing_conn_details) = ctx.connection_details {
@@ -104,31 +116,19 @@ impl IcebergContextManager {
             }
         }
 
-        // cancel any in-progress action and connect to the new connection
         {
             let mut ctx = self.iceberg_context.write().unwrap();
-            // TODO: cancel in-prog action
-            // if let Some(cancellable) = *ctx.deref_mut().cancellable_action {
-            //     cancellable.abort();
-            // }
             ctx.connect_to(new_conn_details);
         }
 
-        // spawn a task to start populating the namespaces
-        let action_tx = self.action_tx.clone();
-        let ctx = self.iceberg_context.clone();
-        // TODO: store the join handle for cancellation
-        let _jh = tokio::spawn(async move {
-            let res = Self::populate_namespaces(ctx, action_tx).await;
-            if let Err(error) = res {
-                tracing::error!(%error, "Error populating namespaces");
-            }
-        });
-
         Ok(())
     }
 
-    async fn populate_namespaces(ctx: IceCtxRef, action_tx: ActionTx) -> Result<()> {
+    async fn populate_namespaces(
+        ctx: IceCtxRef,
+        action_tx: ActionTx,
+        job_queue_tx: MpscSender<IcebergTask>,
+    ) -> Result<()> {
         let root_namespaces = {
             let catalog = {
                 let r_ctx = ctx.read().unwrap();
@@ -157,23 +157,30 @@ impl IcebergContextManager {
 
         action_tx
             .send(TanicAction::UpdateNamespacesList(
-                namespaces.iter().map(|ns| {
-                    ns.name.clone()
-                }).collect::<Vec<_>>()
+                namespaces
+                    .iter()
+                    .map(|ns| ns.name.clone())
+                    .collect::<Vec<_>>(),
             ))
             .map_err(|err| TanicError::UnexpectedError(err.to_string()))?;
 
+        for namespace in namespaces {
+            let _ = job_queue_tx
+                .send(IcebergTask::TablesForNamespace(namespace.clone()))
+                .await;
+        }
+
         Ok(())
     }
 
     async fn populate_tables(
         ctx: IceCtxRef,
         action_tx: ActionTx,
         namespace: NamespaceDeets,
+        job_queue_tx: MpscSender<IcebergTask>,
     ) -> Result<()> {
         let namespace_ident = NamespaceIdent::from_strs(namespace.parts.clone())?;
         let tables = {
-
             let catalog = {
                 let r_ctx = ctx.read().unwrap();
 
@@ -206,10 +213,59 @@ impl IcebergContextManager {
         action_tx
             .send(TanicAction::UpdateNamespaceTableList(
                 namespace.name.clone(),
-                tables.iter().map(|t|&t.name).cloned().collect(),
+                tables.iter().map(|t| &t.name).cloned().collect(),
             ))
             .map_err(TanicError::unexpected)?;
 
+        for table in tables {
+            let _ = job_queue_tx
+                .send(IcebergTask::SummaryForTable(table.clone()))
+                .await;
+        }
+
+        Ok(())
+    }
+
+    async fn populate_table_summary(
+        ctx: IceCtxRef,
+        action_tx: ActionTx,
+        table: TableDeets,
+        _job_queue_tx: MpscSender<IcebergTask>,
+    ) -> Result<()> {
+        let namespace_ident = NamespaceIdent::from_strs(table.namespace.clone())?;
+        let table_ident = TableIdent::new(namespace_ident.clone(), table.name.clone());
+
+        let loaded_table = {
+            let catalog = {
+                let r_ctx = ctx.read().unwrap();
+
+                let Some(ref catalog) = r_ctx.catalog else {
+                    return Err(TanicError::unexpected(
+                        "Attempted to populate table summary when catalog not initialised",
+                    ));
+                };
+
+                catalog.clone()
+            };
+
+            catalog.load_table(&table_ident).await?
+        };
+
+        let summary = loaded_table
+            .metadata()
+            .current_snapshot()
+            .unwrap()
+            .summary();
+        tracing::info!(?summary);
+
+        action_tx
+            .send(TanicAction::UpdateTableSummary {
+                namespace: namespace_ident.to_url_string(),
+                table_name: table_ident.name.clone(),
+                table_summary: summary.additional_properties.clone(),
+            })
+            .map_err(TanicError::unexpected)?;
+
         Ok(())
     }
 }
@@ -229,3 +285,62 @@ impl IcebergContext {
         self.tables = vec![];
     }
 }
+
+impl IcebergContextManager {
+    async fn job_handler(
+        job_queue_rx: MpscReceiver<IcebergTask>,
+        job_queue_tx: MpscSender<IcebergTask>,
+        action_tx: ActionTx,
+        iceberg_ctx: IceCtxRef,
+    ) {
+        let job_stream = ReceiverStream::new(job_queue_rx);
+
+        // let _ = tokio::spawn(async move {
+        job_stream
+            .map(|task| {
+                (
+                    task,
+                    iceberg_ctx.clone(),
+                    action_tx.clone(),
+                    job_queue_tx.clone(),
+                )
+            })
+            .for_each_concurrent(
+                JOB_STREAM_CONCURRENCY,
+                async move |(task, iceberg_ctx, action_tx, job_queue_tx)| {
+                    match task {
+                        IcebergTask::Namespaces => {
+                            let _ = IcebergContextManager::populate_namespaces(
+                                iceberg_ctx,
+                                action_tx,
+                                job_queue_tx,
+                            )
+                            .await;
+                        }
+
+                        IcebergTask::TablesForNamespace(namespace) => {
+                            let _ = IcebergContextManager::populate_tables(
+                                iceberg_ctx,
+                                action_tx,
+                                namespace,
+                                job_queue_tx,
+                            )
+                            .await;
+                        }
+
+                        IcebergTask::SummaryForTable(table) => {
+                            let _ = IcebergContextManager::populate_table_summary(
+                                iceberg_ctx,
+                                action_tx,
+                                table,
+                                job_queue_tx,
+                            )
+                            .await;
+                        } // _ => {}
+                    }
+                },
+            )
+            .await;
+        // }).await;
+    }
+}

tanic-svc/src/lib.rs

@@ -8,8 +8,8 @@ use tokio::sync::watch::{Receiver as WatchReceiver, Sender as WatchSender};
 pub mod iceberg_context;
 pub mod state;
 
-pub use state::{TanicAction, TanicAppState};
 use crate::state::TanicIcebergState;
+pub use state::{TanicAction, TanicAppState};
 
 pub struct AppStateManager {
     action_rx: MpscReceiver<TanicAction>,
@@ -22,9 +22,7 @@ pub struct AppStateManager {
 }
 
 impl AppStateManager {
-    pub fn new(
-        _config: TanicConfig,
-    ) -> (Self, MpscSender<TanicAction>, WatchReceiver<()>) {
+    pub fn new(_config: TanicConfig) -> (Self, MpscSender<TanicAction>, WatchReceiver<()>) {
         let state = Arc::new(RwLock::new(TanicAppState::default()));
 
         let (action_tx, action_rx) = tokio::sync::mpsc::unbounded_channel();