Merge pull request #248 from fMeow/main

feat: add `timestamp-query` feature to allow timestamp offset query

Author: crepererum

src/client/consumer.rs

@@ -502,6 +502,9 @@ mod tests {
                 match at {
                     OffsetAt::Earliest => Ok(inner.lock().await.range.0),
                     OffsetAt::Latest => Ok(inner.lock().await.range.1),
+                    OffsetAt::Timestamp(_) => {
+                        unreachable!("timestamp based offset is tested in e2e test")
+                    }
                 }
             })
         }

src/client/partition.rs

@@ -23,7 +23,7 @@ use crate::{
     throttle::maybe_throttle,
     validation::ExactlyOne,
 };
-use chrono::{LocalResult, TimeZone, Utc};
+use chrono::{DateTime, LocalResult, TimeZone, Utc};
 use std::{
     ops::{ControlFlow, Deref, Range},
     sync::Arc,
@@ -82,11 +82,6 @@ pub enum Compression {
 }
 
 /// Which type of offset should be requested by [`PartitionClient::get_offset`].
-///
-/// # Timestamp-based Queries
-/// In theory the Kafka API would also support querying an offset based on a timestamp, but the behavior seems to be
-/// semi-defined, unintuitive (even within Apache Kafka) and inconsistent between Apache Kafka and Redpanda. So we
-/// decided to NOT expose this option.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum OffsetAt {
     /// Earliest existing record.
@@ -97,6 +92,13 @@ pub enum OffsetAt {
 
     /// The latest existing record.
     Latest,
+
+    /// Timestamp
+    ///
+    /// The Kafka API supports querying an offset based on a timestamp, but the behavior seems to be
+    /// semi-defined, unintuitive (even within Apache Kafka) and inconsistent between Apache Kafka and Redpanda. Mostly it
+    /// offers timestamp based querying with millisecond precision.
+    Timestamp(DateTime<Utc>),
 }
 
 #[derive(Debug)]
@@ -902,6 +904,7 @@ fn build_list_offsets_request(partition: i32, topic: &str, at: OffsetAt) -> List
     let timestamp = match at {
         OffsetAt::Earliest => -2,
         OffsetAt::Latest => -1,
+        OffsetAt::Timestamp(ts) => ts.timestamp_millis(),
     };
 
     ListOffsetsRequest {

tests/consumer.rs

@@ -2,9 +2,12 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use assert_matches::assert_matches;
+use chrono::{DateTime, TimeZone, Utc};
 use futures::{Stream, StreamExt};
 use tokio::time::timeout;
 
+use rskafka::client::partition::OffsetAt;
+use rskafka::record::Record;
 use rskafka::{
     client::{
         consumer::{StartOffset, StreamConsumer, StreamConsumerBuilder},
@@ -415,6 +418,60 @@ async fn test_stream_consumer_start_at_latest_empty() {
     assert_stream_pending(&mut stream).await;
 }
 
+#[tokio::test]
+async fn test_stream_consumer_start_timestamp_based_offset() {
+    maybe_start_logging();
+
+    let test_cfg = maybe_skip_kafka_integration!();
+    let client = ClientBuilder::new(test_cfg.bootstrap_brokers)
+        .build()
+        .await
+        .unwrap();
+    let controller_client = client.controller_client().unwrap();
+
+    let topic = random_topic_name();
+    controller_client
+        .create_topic(&topic, 1, 1, 5_000)
+        .await
+        .unwrap();
+
+    let partition_client = Arc::new(
+        client
+            .partition_client(&topic, 0, UnknownTopicHandling::Retry)
+            .await
+            .unwrap(),
+    );
+    let ts = Utc.timestamp_millis_opt(1337).unwrap();
+    let record_1 = record_with_timestamp_milliseconds(b"x", ts);
+    let record_2 = record_with_timestamp_milliseconds(b"y", ts + Duration::from_millis(100));
+    partition_client
+        .produce(vec![record_1.clone()], Compression::NoCompression)
+        .await
+        .unwrap();
+
+    partition_client
+        .produce(vec![record_2.clone()], Compression::NoCompression)
+        .await
+        .unwrap();
+
+    let offset = partition_client
+        .get_offset(OffsetAt::Timestamp(ts + Duration::from_millis(100)))
+        .await
+        .unwrap();
+    assert_eq!(offset, 1);
+    let mut stream =
+        StreamConsumerBuilder::new(Arc::clone(&partition_client), StartOffset::At(offset))
+            .with_max_wait_ms(50)
+            .build();
+
+    // Get record
+    let (record_and_offset, _) = assert_ok(timeout(TEST_TIMEOUT, stream.next()).await);
+    assert_eq!(record_and_offset.record, record_2);
+
+    // No further records
+    assert_stream_pending(&mut stream).await;
+}
+
 fn assert_ok(
     r: Result<Option<<StreamConsumer as Stream>::Item>, tokio::time::error::Elapsed>,
 ) -> (RecordAndOffset, i64) {
@@ -436,3 +493,12 @@ where
         _ = tokio::time::sleep(Duration::from_secs(1)) => {},
     };
 }
+
+fn record_with_timestamp_milliseconds(key: &[u8], timestamp: DateTime<Utc>) -> Record {
+    Record {
+        key: Some(key.to_vec()),
+        value: Some(b"hello kafka".to_vec()),
+        headers: std::collections::BTreeMap::from([("foo".to_owned(), b"bar".to_vec())]),
+        timestamp,
+    }
+}