From 22636f43f20e7cfece9fbf75d48cb2ff0b71e18d Mon Sep 17 00:00:00 2001
From: Joseph Shearer <joseph@josephshearer.net>
Date: Fri, 18 Oct 2024 12:27:16 -0400
Subject: [PATCH] dekaf: Always return at least 2 docs for data preview UI
 requests

---
 crates/dekaf/src/session.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/crates/dekaf/src/session.rs b/crates/dekaf/src/session.rs
index d2e6b4daf4..4a33e2c853 100644
--- a/crates/dekaf/src/session.rs
+++ b/crates/dekaf/src/session.rs
@@ -23,7 +23,7 @@ use std::{
     collections::{hash_map::Entry, HashMap},
     time::{SystemTime, UNIX_EPOCH},
 };
-use std::{sync::Arc, time::Duration};
+use std::{sync::Arc, time::Duration, cmp::max};
 use tracing::instrument;
 
 struct PendingRead {
@@ -491,7 +491,10 @@ impl Session {
                                     Some(partition_request.fetch_offset - 1),
                                 )
                                 .next_batch(
-                                    crate::read::ReadTarget::Docs(diff as usize + 1),
+                                    // Have to read at least 2 docs, as the very last doc
+                                    // will probably be a control document and will be
+                                    // ignored by the consumer, looking like 0 docs were read
+                                    crate::read::ReadTarget::Docs(max(diff as usize, 2)),
                                     std::time::Instant::now() + timeout,
                                 ),
                             )