🔥 Realtime Data Enrichment - add get_enrichment_table_record fn to VRL log transform pipeline

Author: shaeqahmed

example/enrichment/user_info/data.json

@@ -0,0 +1 @@
+{"user_id":"john","name":"John Doe"}

example/enrichment/user_info/enrichment.yml

@@ -0,0 +1,13 @@
+name: "user_info"
+
+enrichment_type: "static"
+
+# Iceberg schema (same as log sources)
+schema:
+  fields:
+    - name: name
+      type: string
+    - name: user_id
+      type: string
+lookup_keys:
+  - user_id

infra/src/DPMainStack.ts

@@ -199,6 +199,14 @@ export class DPMainStack extends MatanoStack {
     });
     transformer.node.addDependency(sqsSources);
 
+    if (enrichment != null) {
+      transformer.transformerLambda.addEnvironment(
+        "ENRICHMENT_TABLES_BUCKET",
+        enrichment.enrichmentTablesBucket.bucketName
+      );
+      enrichment.enrichmentTablesBucket.grantRead(transformer.transformerLambda);
+    }
+
     const rawDataBatcher = new DataBatcher(this, "DataBatcher", {
       transformerFunction: transformer.transformerLambda,
       s3Bucket: props.matanoSourcesBucket,

lib/java/matano/iceberg_main/src/main/kotlin/com/matano/iceberg/AthenaUtil.kt

@@ -1,12 +1,14 @@
 package com.matano.iceberg
 
+import org.slf4j.LoggerFactory
 import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient
 import software.amazon.awssdk.services.athena.AthenaAsyncClient
 import software.amazon.awssdk.services.athena.model.*
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.TimeUnit
 
 class AthenaQueryRunner(val workGroup: String = "matano_system") {
+    private val logger = LoggerFactory.getLogger(this::class.java)
     val SLEEP_AMOUNT_MS = 500L
     val delayedExecutor = CompletableFuture.delayedExecutor(SLEEP_AMOUNT_MS, TimeUnit.MILLISECONDS)
 
@@ -38,9 +40,12 @@ class AthenaQueryRunner(val workGroup: String = "matano_system") {
                     )
                 }
                 QueryExecutionState.CANCELLED -> throw RuntimeException("The Amazon Athena query was cancelled.")
-                QueryExecutionState.SUCCEEDED -> CompletableFuture.completedFuture(Unit)
+                QueryExecutionState.SUCCEEDED -> {
+                    logger.info("Successfully completed Athena query.")
+                    CompletableFuture.completedFuture(Unit)
+                }
                 else -> {
-                    CompletableFuture.supplyAsync({ waitForQueryToComplete(queryExecutionId) }, delayedExecutor)
+                    CompletableFuture.supplyAsync({}, delayedExecutor).thenComposeAsync { waitForQueryToComplete(queryExecutionId) }
                 }
             }
         }.thenApplyAsync {}

lib/java/matano/iceberg_main/src/main/kotlin/com/matano/iceberg/Enrichment.kt

@@ -157,7 +157,7 @@ class EnrichmentIcebergSyncer {
             TO 's3://$tempSyncBucket/$keyPrefix'
             WITH (format = 'AVRO')
         """.trimIndent()
-        athenaQueryRunner.runAthenaQuery(qs)
+        athenaQueryRunner.runAthenaQuery(qs).await()
 
         // Make downloads concurrent
         val futs = s3AsyncClient.listObjectsV2 { r -> r.bucket(tempSyncBucket).prefix(keyPrefix) }.await().contents().map { res ->
@@ -229,16 +229,18 @@ class EnrichmentIcebergSyncer {
             TO 's3://$tempSyncBucket/$keyPrefix'
             WITH (format = 'PARQUET', compression='snappy')
         """.trimIndent()
-        athenaQueryRunner.runAthenaQuery(qs)
+        athenaQueryRunner.runAthenaQuery(qs).await()
 
-        val futs = s3AsyncClient.listObjectsV2 { r -> r.bucket(tempSyncBucket).prefix(keyPrefix) }.await().contents().map { res ->
+        val ll = s3AsyncClient.listObjectsV2 { r -> r.bucket(tempSyncBucket).prefix(keyPrefix) }.await().contents().toList()
+        val futs = ll.map { res ->
             s3AsyncClient
                 .getObject({ it.bucket(tempSyncBucket).key(res.key()) }, AsyncResponseTransformer.toBytes())
                 .thenApply { InMemoryInputFile(it.asByteArray()) }
         }
         CompletableFuture.allOf(*futs.toTypedArray()).await()
         val inputFiles = futs.map { it.get() }
         if (inputFiles.isEmpty()) {
+            logger.info("Input files are empty.")
             return null
         }
 

lib/rust/shared/src/avro_index.rs

@@ -1,26 +1,27 @@
+use crate::avro::AvroValueExt;
 use apache_avro::from_avro_datum;
+use log::info;
 use memmap2::{Mmap, MmapOptions};
 use once_cell::sync::OnceCell;
-use crate::avro::AvroValueExt;
 use std::io::{BufRead, Read};
 use std::io::{BufReader, Seek};
 use std::vec;
 use std::{
     collections::HashMap,
     io::Cursor,
-    sync::{Arc, Mutex},
+    sync::{Arc, Mutex, Weak},
 };
 
 use anyhow::{anyhow, Context, Result};
 
 /// We just need this because the Avro reader takes ownership but we want to be able to use the underlying reader to seek.
 struct ReaderHolder {
-    reader: Arc<Mutex<Cursor<Mmap>>>,
+    reader: Weak<Mutex<Cursor<Mmap>>>,
 }
 
 impl std::io::Read for ReaderHolder {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-        self.reader.lock().unwrap().read(buf)
+        self.reader.upgrade().unwrap().lock().unwrap().read(buf)
     }
 }
 
@@ -41,12 +42,12 @@ type RecordPosMap = HashMap<String, [i64; 2]>; // [start_pos, block_length]
 pub struct AvroIndex {
     pub schema: apache_avro::Schema,
     indices: HashMap<String, (String, OnceCell<RecordPosMap>)>,
-    reader: Arc<Mutex<Cursor<Mmap>>>,
+    reader: Arc<Mmap>,
 }
 
 impl AvroIndex {
-    /// There's an offset of 3 bytes between the start positions written in Index and where raw compressed block starts.
-    const OFFSET: u64 = 3;
+    /// There's an offset of 2 bytes between the start positions written in Index and where raw compressed block starts.
+    const OFFSET: u64 = 2;
     /// The length of the sync marker at the end of each block. Remove to get actual data.
     const SYNC_LENGTH: usize = 16;
 
@@ -63,15 +64,24 @@ impl AvroIndex {
         let reader = Arc::new(Mutex::new(Cursor::new(mmap)));
 
         let reader_holder = ReaderHolder {
-            reader: reader.clone(),
+            reader: Arc::downgrade(&reader.clone()),
         };
 
         let avro_reader = apache_avro::Reader::new(reader_holder)?;
+        let schema = avro_reader.writer_schema().clone();
+
+        let mem_cursor = Arc::try_unwrap(reader)
+            .map_err(|e| anyhow!("Could not unwrap Arc: {:?}", e))?
+            .into_inner()?;
+
+        let mem = mem_cursor.into_inner();
+
+        let reader = Arc::new(mem);
 
         Ok(Self {
             indices,
             reader,
-            schema: avro_reader.writer_schema().clone(),
+            schema,
         })
     }
 
@@ -120,19 +130,17 @@ impl AvroIndex {
                 let block_len = *block_len as usize - (Self::OFFSET as usize + Self::SYNC_LENGTH);
 
                 // Read the whole block since we know the length and avoid multiple reads.
-                // Need scope to ensure mutex is dropped and released before avro reads.
                 let raw_block: Vec<u8> = {
-                    let mut reader = self.reader.lock().unwrap();
-                    reader.seek(std::io::SeekFrom::Start(start_pos))?;
-
-                    let mut buf = vec![0; block_len as usize];
-                    reader.read_exact(buf.as_mut())?;
-
+                    let reader = self.reader.clone();
+                    let start_pos = start_pos as usize;
+                    let buf = reader[start_pos..start_pos + block_len].to_vec();
                     buf
                 };
 
                 let mut block = vec![];
-                zstd::Decoder::new(raw_block.as_slice())?.read_to_end(block.as_mut())?;
+                zstd::Decoder::new(raw_block.as_slice())?
+                    .read_to_end(block.as_mut())
+                    .map_err(|e| anyhow!("Failed to decompress avro idx block").context(e))?;
 
                 let mut block_reader = block.as_slice();
                 while !block_reader.is_empty() {

lib/rust/shared/src/enrichment.rs

@@ -15,6 +15,7 @@ use async_compression::tokio::bufread::ZstdDecoder;
 use pyo3::prelude::*;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::runtime::Runtime;
+use tracing::log::info;
 
 use crate::avro_index::AvroIndex;
 use crate::utils::load_enrichment_config;
@@ -88,6 +89,12 @@ pub fn create_enrichment_tables(table_names: Vec<String>) -> PyResult<Vec<Enrich
     Ok(tables)
 }
 
+pub fn create_enrichment_table(table_name: &str) -> Result<EnrichmentTable> {
+    let lookup_keys = get_enrichment_table_lookup_keys(table_name)
+        .ok_or_else(|| anyhow!("No lookup key(s) for table {}", table_name))?;
+    Ok(EnrichmentTable::new(table_name.to_string(), lookup_keys))
+}
+
 fn get_enrichment_table_lookup_keys(table_name: &str) -> Option<Vec<String>> {
     lazy_static! {
         static ref ENRICHMENT_CONFIG: HashMap<String, serde_yaml::Value> =
@@ -109,6 +116,7 @@ async fn load_avro_index(table_name: &str, lookup_keys: &Vec<String>) -> Result<
     let bucket = std::env::var("ENRICHMENT_TABLES_BUCKET")?;
     let avro_key = format!("tables/{}.avro", table_name);
     let local_dir = std::env::temp_dir().join("tables");
+    tokio::fs::create_dir_all(&local_dir).await?;
     let local_avro_path = local_dir.join(format!("{}.zstd.avro", table_name));
 
     let load_avro = || async {
@@ -122,12 +130,11 @@ async fn load_avro_index(table_name: &str, lookup_keys: &Vec<String>) -> Result<
             .into_async_read();
         let mut avro_reader = tokio::io::BufReader::new(avro_reader);
 
-        tokio::fs::create_dir_all(&local_dir).await?;
-
         let file = tokio::fs::File::create(&local_avro_path).await?;
         let mut writer = tokio::io::BufWriter::new(file);
 
         tokio::io::copy(&mut avro_reader, &mut writer).await?;
+        
         anyhow::Ok(HashMap::<String, String>::with_capacity(0)) // just to align types
     };
 
@@ -153,10 +160,10 @@ async fn load_avro_index(table_name: &str, lookup_keys: &Vec<String>) -> Result<
                     let mut index_reader =
                         ZstdDecoder::new(tokio::io::BufReader::new(index_reader));
 
-                    let file = tokio::fs::File::create(&local_index_path).await?;
+                    let file = tokio::fs::File::create(&local_index_path).await.map_err(|e| anyhow!("failed to create file {}", local_index_path).context(e))?;
                     let mut writer = tokio::io::BufWriter::new(file);
 
-                    tokio::io::copy(&mut index_reader, &mut writer).await?;
+                    tokio::io::copy(&mut index_reader, &mut writer).await.map_err(|e| anyhow!("failed to copy index").context(e))?;
                     anyhow::Ok((lookup_key.to_string(), local_index_path))
                 }
             })