ydb 0.17.1

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

use std::time::Duration;

use tracing::instrument;

use crate::{
    TopicReader, TopicReaderBatch, YdbError, YdbResult,
    client_query::{
        QueryTxIdentity, Transaction,
        hooks::{QueryTxCommitStatus, QueryTxHook},
    },
    grpc_wrapper::{
        raw_topic_service::{
            client::RawTopicClient,
            common::partition::RawOffsetsRange,
            update_offsets_in_transaction::{
                RawPartitionOffsets, RawTopicOffsets, RawTransactionIdentity,
                RawUpdateOffsetsInTransactionRequest,
            },
        },
        raw_ydb_operation::RawOperationParams,
    },
};

use super::runtime::RuntimeHandle;

struct ReaderTxHook {
    runtime: RuntimeHandle,
}

impl QueryTxHook for ReaderTxHook {
    fn after_commit(&mut self, status: QueryTxCommitStatus) {
        match status {
            QueryTxCommitStatus::Committed => {}

            QueryTxCommitStatus::Aborted => {
                let _ = self.runtime.force_reconnection(YdbError::custom(
                    "query transaction was aborted: force new connection",
                ));
            }
        }
    }
}

pub struct TopicReaderTx<'a> {
    inner: &'a mut TopicReader,
    runtime: RuntimeHandle,
    client: RawTopicClient,
    tx_identity: QueryTxIdentity,
}

impl<'a> TopicReaderTx<'a> {
    pub(super) async fn new(reader: &'a mut TopicReader, tx: &mut Transaction) -> YdbResult<Self> {
        let uri = tx.uri().await?.ok_or(YdbError::custom("no node Uri"))?;

        let client = reader
            .manager
            .get_auth_service_to_node(RawTopicClient::new, uri)
            .await?;

        let tx_identity = tx.tx_identity().await?;
        let runtime = reader.runtime_handle();

        tx.register_hook(ReaderTxHook {
            runtime: runtime.clone(),
        });

        Ok(Self {
            inner: reader,
            runtime,
            client,
            tx_identity,
        })
    }

    #[instrument(name = "ydb.TopicReaderTx.ReadBatch", skip_all, fields(db.system.name = "ydb"), err)]
    pub async fn read_batch(&mut self) -> YdbResult<TopicReaderBatch> {
        let batch = self.inner.read_batch_inner().await?;
        if let Err(err) = self.update_offsets_in_transaction(&batch).await {
            let _ = self.runtime.force_reconnection(YdbError::custom(
                "UpdateOffsetsInTransaction failed after reading batch",
            ));
            return Err(err);
        }
        Ok(batch)
    }

    async fn update_offsets_in_transaction(&mut self, batch: &TopicReaderBatch) -> YdbResult<()> {
        let commit_marker = batch.get_commit_marker();

        let request = RawUpdateOffsetsInTransactionRequest {
            operation_params: RawOperationParams::new_with_timeouts(
                Duration::from_secs(30),
                Duration::from_secs(60),
            ),
            tx: RawTransactionIdentity {
                id: self.tx_identity.transaction_id.clone(),
                session: self.tx_identity.session_id.clone(),
            },
            topics: vec![RawTopicOffsets {
                path: commit_marker.topic.clone(),
                partitions: vec![RawPartitionOffsets {
                    partition_id: commit_marker.partition_id.into_raw(),
                    partition_offsets: vec![RawOffsetsRange {
                        start: commit_marker.start_offset,
                        end: commit_marker.end_offset,
                    }],
                }],
            }],
            consumer: self.inner.consumer().to_string(),
        };

        self.client.update_offsets_in_transaction(request).await?;

        Ok(())
    }
}