ref(api): Improve performance of the replication status endpoint (#484)

Author: iambriccardo

etl-api/src/routes/pipelines.rs

@@ -4,7 +4,9 @@ use actix_web::{
     post,
     web::{Data, Json, Path},
 };
-use etl_postgres::replication::{TableLookupError, get_table_name_from_oid, health, lag, state};
+use etl_postgres::replication::{
+    TableLookupError, get_table_names_from_table_ids, health, lag, state,
+};
 use etl_postgres::types::TableId;
 use serde::{Deserialize, Serialize};
 use sqlx::PgPool;
@@ -939,7 +941,7 @@ pub async fn get_pipeline_replication_status(
 
     txn.commit().await?;
 
-    // Connect to the source database to read replication state
+    // Connect to the source database to read the necessary state
     let source_pool =
         connect_to_source_database_with_defaults(&source.config.into_connection_config()).await?;
 
@@ -953,11 +955,20 @@ pub async fn get_pipeline_replication_status(
     let mut lag_metrics = lag::get_pipeline_lag_metrics(&source_pool, pipeline_id as u64).await?;
     let apply_lag = lag_metrics.apply.map(Into::into);
 
-    // Convert database states to UI-friendly format and fetch table names
+    // Collect all table IDs and fetch their names in a single batch query
+    let table_ids: Vec<TableId> = state_rows
+        .iter()
+        .map(|row| TableId::new(row.table_id.0))
+        .collect();
+    let table_names = get_table_names_from_table_ids(&source_pool, &table_ids).await?;
+
+    // Convert database states to UI-friendly format
     let mut tables: Vec<TableReplicationStatus> = Vec::new();
     for row in state_rows {
         let table_id = TableId::new(row.table_id.0);
-        let table_name = get_table_name_from_oid(&source_pool, table_id).await?;
+        let table_name = table_names
+            .get(&table_id)
+            .ok_or(TableLookupError::TableNotFound(table_id))?;
 
         // Extract the metadata row from the database
         let table_replication_state = row

etl-postgres/src/replication/db.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::num::NonZeroI32;
 
 use etl_config::shared::{IntoConnectOptions, PgConnectionConfig};
@@ -37,38 +38,44 @@ pub async fn connect_to_source_database(
     Ok(pool)
 }
 
-/// Retrieves table name from table OID by querying system catalogs.
+/// Retrieves table names for multiple table ids in a single query.
 ///
-/// Looks up the schema and table name for the given table OID using Postgres's
-/// pg_class and pg_namespace system tables.
-pub async fn get_table_name_from_oid(
+/// Looks up the schema and table names for all given table OIDs using Postgres's
+/// pg_class and pg_namespace system tables. Returns a HashMap mapping each TableId
+/// to its corresponding TableName.
+pub async fn get_table_names_from_table_ids(
     pool: &PgPool,
-    table_id: TableId,
-) -> Result<TableName, TableLookupError> {
-    let query = "
-        select n.nspname as schema_name, c.relname as table_name
-        from pg_class c
-        join pg_namespace n on c.relnamespace = n.oid
-        where c.oid = $1
+    table_ids: &[TableId],
+) -> Result<HashMap<TableId, TableName>, TableLookupError> {
+    if table_ids.is_empty() {
+        return Ok(HashMap::new());
+    }
+
+    let query = "select c.oid::int as oid, n.nspname as schema_name, c.relname as table_name
+    from pg_class c
+    join pg_namespace n on c.relnamespace = n.oid
+    where c.oid = any($1::oid[])
     ";
 
-    let row = sqlx::query(query)
-        .bind(table_id.into_inner() as i64)
-        .fetch_optional(pool)
-        .await?;
+    let oids: Vec<i32> = table_ids.iter().map(|id| id.into_inner() as i32).collect();
+    let rows = sqlx::query(query).bind(&oids).fetch_all(pool).await?;
 
-    match row {
-        Some(row) => {
-            let schema_name: String = row.try_get("schema_name")?;
-            let table_name: String = row.try_get("table_name")?;
+    let mut result = HashMap::with_capacity(rows.len());
+    for row in rows {
+        let oid: i32 = row.try_get("oid")?;
+        let schema_name: String = row.try_get("schema_name")?;
+        let table_name: String = row.try_get("table_name")?;
 
-            Ok(TableName {
+        result.insert(
+            TableId::new(oid as u32),
+            TableName {
                 schema: schema_name,
                 name: table_name,
-            })
-        }
-        None => Err(TableLookupError::TableNotFound(table_id)),
+            },
+        );
     }
+
+    Ok(result)
 }
 
 /// Extracts the PostgreSQL server version from a version string.