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

use tracing::error;
use ydb_grpc::ydb_proto::query::{
    BeginTransactionResponse, CommitTransactionResponse, CreateSessionResponse,
    DeleteSessionResponse, RollbackTransactionResponse, SessionState, TransactionMeta,
};
use ydb_grpc::ydb_proto::status_ids::StatusCode;

use crate::mock_server::handler::{FromHandlerToService, Handler, Incoming, Reply};

use super::handler::{QueryIncoming, QueryReply};

pub const QUERY_SESSION_ID: &str = "session-id-xyz";
pub const QUERY_TX_ID: &str = "tx-id-abc";

pub struct QueryDefaultHandler {
    tx: FromHandlerToService,
}

impl QueryDefaultHandler {
    pub fn with_tx(tx: FromHandlerToService) -> Self {
        Self { tx }
    }
}

impl Handler for QueryDefaultHandler {
    fn set_channel(&mut self, tx: FromHandlerToService) {
        self.tx = tx;
    }

    fn handle(&self, incoming: Incoming) -> Option<Incoming> {
        let Incoming::Query(incoming) = incoming else {
            error!(?incoming, "query default handler got non-query message");
            return Some(incoming);
        };

        match incoming {
            QueryIncoming::CreateSession(_, reply_tx) => {
                let _ = reply_tx.send(Ok(tonic::Response::new(CreateSessionResponse {
                    status: StatusCode::Success as i32,
                    issues: Vec::new(),
                    session_id: QUERY_SESSION_ID.to_string(),
                    node_id: 0,
                })));
            }
            QueryIncoming::DeleteSession(_, reply_tx) => {
                let _ = reply_tx.send(Ok(tonic::Response::new(DeleteSessionResponse {
                    status: StatusCode::Success as i32,
                    issues: Vec::new(),
                })));
            }
            QueryIncoming::BeginTransaction(_, reply_tx) => {
                let _ = reply_tx.send(Ok(tonic::Response::new(BeginTransactionResponse {
                    status: StatusCode::Success as i32,
                    issues: Vec::new(),
                    tx_meta: Some(TransactionMeta {
                        id: QUERY_TX_ID.to_string(),
                    }),
                })));
            }
            QueryIncoming::CommitTransaction(_, reply_tx) => {
                let _ = reply_tx.send(Ok(tonic::Response::new(CommitTransactionResponse {
                    status: StatusCode::Success as i32,
                    issues: Vec::new(),
                })));
            }
            QueryIncoming::RollbackTransaction(_, reply_tx) => {
                let _ = reply_tx.send(Ok(tonic::Response::new(RollbackTransactionResponse {
                    status: StatusCode::Success as i32,
                    issues: Vec::new(),
                })));
            }
            QueryIncoming::AttachSession(_, stream_id) => {
                let _ = self.tx.send(Reply::Query(QueryReply::AttachSession {
                    stream_id,
                    state: SessionState {
                        status: StatusCode::Success as i32,
                        issues: Vec::new(),
                        session_hint: None,
                    },
                }));
                // stream stays open — no close sent, mirrors stream::pending() behavior
            }
            QueryIncoming::ExecuteQuery(_, stream_id) => {
                let _ = self
                    .tx
                    .send(Reply::Query(QueryReply::ExecuteQueryClose { stream_id }));
            }
        }

        None
    }
}