Add more tracing and instrumentation

Signed-off-by: Abhi Agarwal <[email protected]>

Author: abhi-airspace-intelligence

etl-destinations/src/deltalake/maintenance.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use deltalake::DeltaTable;
 use tokio::sync::Mutex;
 use tokio::task::JoinHandle;
-use tracing::{error, trace};
+use tracing::{debug, error, instrument, trace};
 
 use etl::types::TableId;
 
@@ -38,6 +38,7 @@ impl TableMaintenanceState {
 
     /// Await any in-flight compaction, then if the `compact_after_commits` threshold is met,
     /// run compaction. This guarantees serialization of compaction runs relative to table writes.
+    #[instrument(skip(self, table, config), fields(table_id = table_id.0, table_version))]
     pub async fn maybe_run_compaction(
         self: &Arc<Self>,
         table_id: TableId,
@@ -69,6 +70,11 @@ impl TableMaintenanceState {
         };
 
         if !should_compact {
+            debug!(
+                table_id = table_id.0,
+                version = table_version,
+                "Skipping Delta table compaction task"
+            );
             return;
         }
 
@@ -103,6 +109,7 @@ impl TableMaintenanceState {
 
     /// Await any in-flight Z-ordering, then if the `z_order_after_commits` threshold is met,
     /// run Z-order. Serializes Z-order runs relative to table writes.
+    #[instrument(skip(self, table, config), fields(table_id = table_id.0, table_version))]
     pub async fn maybe_run_zorder(
         self: &Arc<Self>,
         table_id: TableId,
@@ -140,6 +147,11 @@ impl TableMaintenanceState {
         };
 
         if !should_zorder {
+            debug!(
+                table_id = table_id.0,
+                version = table_version,
+                "Skipping Delta table Z-order task"
+            );
             return;
         }
 

etl-destinations/src/deltalake/operations/append.rs

@@ -3,18 +3,23 @@ use deltalake::{
     arrow::array::RecordBatch,
     writer::{DeltaWriter, RecordBatchWriter},
 };
+use tracing::{instrument, trace};
 
 use crate::deltalake::config::DeltaTableConfig;
 
 /// Appends a record batch to a Delta table
+#[instrument(skip(table, config, record_batch), fields(num_rows = record_batch.num_rows()))]
 pub async fn append_to_table(
     table: &mut DeltaTable,
     config: &DeltaTableConfig,
     record_batch: RecordBatch,
 ) -> DeltaResult<()> {
+    trace!("Creating RecordBatchWriter for append");
     let mut writer = RecordBatchWriter::for_table(table)?;
     writer = writer.with_writer_properties(config.into());
+    trace!("Writing record batch to Delta table");
     writer.write(record_batch).await?;
+    trace!("Flushing and committing append");
     writer.flush_and_commit(table).await?;
     Ok(())
 }

etl-destinations/src/deltalake/operations/delete.rs

@@ -3,17 +3,22 @@ use deltalake::{
 };
 
 use crate::deltalake::config::DeltaTableConfig;
+use tracing::{instrument, trace};
 
+#[instrument(skip(table, config, delete_predicate))]
 pub async fn delete_from_table(
     table: &mut DeltaTable,
     config: &DeltaTableConfig,
     delete_predicate: Expr,
 ) -> DeltaResult<()> {
+    trace!("Building delete builder with predicate");
     let delete_builder = DeleteBuilder::new((*table).log_store(), table.snapshot()?.clone())
         .with_predicate(delete_predicate)
         .with_writer_properties(config.into());
     // TODO(abhi): Do something with the metrics
+    trace!("Executing delete operation");
     let (deleted_table, _metrics) = delete_builder.await?;
     *table = deleted_table;
+    trace!("Delete operation completed");
     Ok(())
 }

etl-destinations/src/deltalake/operations/merge.rs

@@ -4,6 +4,7 @@ use deltalake::datafusion::prelude::SessionContext;
 use deltalake::operations::merge::MergeBuilder;
 use deltalake::{DeltaResult, DeltaTable, datafusion::prelude::Expr};
 use etl::types::{TableRow as PgTableRow, TableSchema as PgTableSchema};
+use tracing::{instrument, trace};
 
 use crate::arrow::rows_to_record_batch;
 use crate::deltalake::config::DeltaTableConfig;
@@ -14,13 +15,18 @@ pub(crate) fn source_qualified_column_expr(column_name: &str, source_alias: &str
     Expr::Column(Column::new(Some(source_alias), column_name))
 }
 
+#[instrument(
+    skip(table, config, table_schema, upsert_rows, delete_predicate),
+    fields(upsert_count = upsert_rows.len(), has_delete = delete_predicate.is_some())
+)]
 pub async fn merge_to_table(
     table: &mut DeltaTable,
     config: &DeltaTableConfig,
     table_schema: &PgTableSchema,
     upsert_rows: &[PgTableRow],
     delete_predicate: Option<Expr>,
 ) -> DeltaResult<()> {
+    trace!("Building Arrow schema and source batch for merge");
     let arrow_schema = postgres_to_arrow_schema(table_schema)?;
     let rows = rows_to_record_batch(upsert_rows, arrow_schema)?;
 
@@ -38,6 +44,7 @@ pub async fn merge_to_table(
     let qualified_primary_keys = qualify_primary_keys(primary_keys, "source", "target")
         .ok_or(DeltaTableError::generic("Failed to qualify primary keys"))?;
 
+    trace!("Creating merge builder");
     let merge_builder = MergeBuilder::new(
         // TODO(abhi): Is there a way to do this while avoiding the clone/general hackiness?
         (*table).log_store(),
@@ -79,7 +86,9 @@ pub async fn merge_to_table(
             .when_not_matched_by_source_delete(|delete| delete.predicate(delete_predicate))?;
     }
     // TODO(abhi): Do something with the metrics
+    trace!("Executing merge operation");
     let (merged_table, _metrics) = merge_builder.await?;
+    trace!("Merge operation completed");
     *table = merged_table;
     Ok(())
 }

etl-destinations/src/deltalake/operations/optimize.rs

@@ -1,33 +1,40 @@
 use deltalake::operations::optimize::{OptimizeBuilder, OptimizeType};
 use deltalake::parquet::file::properties::WriterProperties;
 use deltalake::{DeltaResult, DeltaTable};
+use tracing::{instrument, trace};
 
 use crate::deltalake::config::DeltaTableConfig;
 
 /// Optimizes a Delta table by compacting small files into larger ones.
+#[instrument(skip(table, config))]
 pub async fn compact_table(table: &mut DeltaTable, config: &DeltaTableConfig) -> DeltaResult<()> {
+    trace!("Starting table compaction");
     let writer_properties = WriterProperties::from(config);
     let optimize_builder = OptimizeBuilder::new(table.log_store(), table.snapshot()?.clone());
     let (optimized_table, _metrics) = optimize_builder
         .with_writer_properties(writer_properties)
         .with_type(OptimizeType::Compact)
         .await?;
     *table = optimized_table;
+    trace!("Finished table compaction");
     Ok(())
 }
 
 /// Optimizes a Delta table by performing Z-order clustering on the provided columns.
+#[instrument(skip(table, config, columns), fields(columns = ?columns))]
 pub async fn zorder_table(
     table: &mut DeltaTable,
     config: &DeltaTableConfig,
     columns: Vec<String>,
 ) -> DeltaResult<()> {
+    trace!("Starting table Z-order optimization");
     let writer_properties = WriterProperties::from(config);
     let optimize_builder = OptimizeBuilder::new(table.log_store(), table.snapshot()?.clone());
     let (optimized_table, _metrics) = optimize_builder
         .with_writer_properties(writer_properties)
         .with_type(OptimizeType::ZOrder(columns))
         .await?;
     *table = optimized_table;
+    trace!("Finished table Z-order optimization");
     Ok(())
 }