ydb 0.10.2

Crate contains generated low-level grpc code from YDB API protobuf, used as base for ydb crate
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

use crate::client_topic::topicreader::partition_state::PartitionSession;
use crate::client_topic::topicreader::reader::TopicReaderCommitMarker;
use crate::grpc_wrapper::raw_topic_service::stream_read::messages::RawBatch;
use crate::YdbResult;
use std::time;
use std::time::SystemTime;

#[cfg_attr(not(feature = "force-exhaustive-all"), non_exhaustive)]
pub struct TopicReaderBatch {
    pub messages: Vec<TopicReaderMessage>,

    commit_marker: TopicReaderCommitMarker,
}

impl TopicReaderBatch {
    pub(crate) fn new(
        raw_batch: RawBatch,
        partition_session: &mut PartitionSession,
    ) -> TopicReaderBatch {
        let written_at: SystemTime = raw_batch.written_at.into();

        let mut batch = Self {
            commit_marker: TopicReaderCommitMarker {
                partition_session_id: partition_session.partition_session_id,
                partition_id: partition_session.partition_id,
                start_offset: partition_session.next_commit_offset_start,
                end_offset: partition_session.next_commit_offset_start,
                topic: partition_session.topic.clone(),
            },

            messages: raw_batch
                .message_data
                .into_iter()
                .map(|message| {
                    let start_commit_offset = partition_session.next_commit_offset_start;
                    partition_session.next_commit_offset_start = message.offset + 1;

                    TopicReaderMessage {
                        seq_no: message.seq_no,
                        created_at: message.created_at.map(|x| x.into()),
                        offset: message.offset,
                        written_at,
                        uncompressed_size: message.uncompressed_size,
                        producer_id: raw_batch.producer_id.clone(),
                        raw_data: Some(message.data),

                        commit_marker: TopicReaderCommitMarker {
                            partition_session_id: partition_session.partition_session_id,
                            partition_id: partition_session.partition_id,
                            start_offset: start_commit_offset,
                            end_offset: message.offset + 1,
                            topic: partition_session.topic.clone(),
                        },
                    }
                })
                .collect(),
        };

        if let Some(last) = batch.messages.last() {
            batch.commit_marker.end_offset = last.commit_marker.end_offset
        }

        batch
    }
}

impl TopicReaderBatch {
    pub fn get_commit_marker(&self) -> TopicReaderCommitMarker {
        self.commit_marker.clone()
    }
}

#[cfg_attr(not(feature = "force-exhaustive-all"), non_exhaustive)]
pub struct TopicReaderMessage {
    pub seq_no: i64,
    pub created_at: Option<time::SystemTime>,
    pub offset: i64,
    pub written_at: time::SystemTime,
    pub uncompressed_size: i64, // as sent by sender, server/sdk doesn't check the field. It may be empty or wrong.

    producer_id: String,
    raw_data: Option<Vec<u8>>,
    commit_marker: TopicReaderCommitMarker,
}

impl TopicReaderMessage {
    pub async fn read_and_take(&mut self) -> YdbResult<Option<Vec<u8>>> {
        Ok(self.raw_data.take())
    }

    pub fn get_producer_id(&self) -> &str {
        self.producer_id.as_str()
    }

    pub fn get_commit_marker(&self) -> TopicReaderCommitMarker {
        self.commit_marker.clone()
    }

    pub fn get_topic(&self) -> &str {
        &self.commit_marker.topic
    }

    pub fn get_partition_id(&self) -> i64 {
        self.commit_marker.partition_id
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::client_topic::topicreader::partition_state::PartitionSession;
    use crate::grpc_wrapper::raw_topic_service::stream_read::messages::{RawBatch, RawMessageData};
    use std::time::SystemTime;

    #[test]
    fn test_topic_reader_batch_new() {
        let mut partition_session = PartitionSession {
            partition_session_id: 123,
            partition_id: 456,
            topic: "test-topic".to_string(),
            next_commit_offset_start: 100,
        };

        let raw_batch = RawBatch {
            producer_id: "test-producer".to_string(),
            write_session_meta: std::collections::HashMap::new(),
            codec: crate::grpc_wrapper::raw_topic_service::common::codecs::RawCodec { code: 1 },
            written_at: SystemTime::now().into(),
            message_data: vec![RawMessageData {
                seq_no: 1,
                created_at: Some(SystemTime::now().into()),
                data: vec![1, 2, 3],
                uncompressed_size: 3,
                offset: 100,
                read_session_size_bytes: 0,
            }],
        };

        let batch = TopicReaderBatch::new(raw_batch, &mut partition_session);

        let commit_marker = batch.get_commit_marker();
        assert_eq!(commit_marker.topic, "test-topic");
        assert_eq!(commit_marker.partition_session_id, 123);
        assert_eq!(commit_marker.partition_id, 456);
        assert_eq!(commit_marker.start_offset, 100);
        assert_eq!(commit_marker.end_offset, 101);

        assert_eq!(batch.messages.len(), 1);
        let message_commit_marker = batch.messages[0].get_commit_marker();
        assert_eq!(message_commit_marker.topic, "test-topic");
        assert_eq!(message_commit_marker.partition_session_id, 123);
        assert_eq!(message_commit_marker.partition_id, 456);
        assert_eq!(message_commit_marker.start_offset, 100);
        assert_eq!(message_commit_marker.end_offset, 101);
    }
}