ydb 0.13.4

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

use std::collections::HashMap;
use std::time::Duration;

use crate::grpc_wrapper::raw_errors::RawResult;
use crate::grpc_wrapper::raw_query_service::status::check_status;
use crate::grpc_wrapper::raw_table_service::value::{
    RawColumn, RawResultSet, RawTypedValue, RawValue,
};
use crate::types::Value;
use ydb_grpc::ydb_proto::query::{
    execute_query_request, ExecMode, ExecuteQueryRequest, ExecuteQueryResponsePart, QueryContent,
    SchemaInclusionMode, StatsMode, Syntax,
};
use ydb_grpc::ydb_proto::result_set::Format;

#[derive(Clone, Debug)]
pub(crate) struct RawExecuteQueryRequest {
    pub session_id: String,
    pub yql_text: String,
    pub parameters: HashMap<String, Value>,
    pub tx_control: Option<ydb_grpc::ydb_proto::query::TransactionControl>,
    pub collect_stats: bool,
}

impl RawExecuteQueryRequest {
    pub(crate) fn new(
        session_id: impl Into<String>,
        yql_text: impl Into<String>,
        parameters: HashMap<String, Value>,
        tx_control: Option<ydb_grpc::ydb_proto::query::TransactionControl>,
        collect_stats: bool,
    ) -> Self {
        Self {
            session_id: session_id.into(),
            yql_text: yql_text.into(),
            parameters,
            tx_control,
            collect_stats,
        }
    }

    pub fn into_proto(self) -> RawResult<ExecuteQueryRequest> {
        let mut parameters = HashMap::with_capacity(self.parameters.len());
        for (name, val) in self.parameters {
            let raw: RawTypedValue = val.try_into()?;
            parameters.insert(name, raw.into());
        }

        Ok(ExecuteQueryRequest {
            session_id: self.session_id,
            exec_mode: ExecMode::Execute as i32,
            tx_control: self.tx_control,
            query: Some(execute_query_request::Query::QueryContent(QueryContent {
                syntax: Syntax::YqlV1 as i32,
                text: self.yql_text,
            })),
            parameters,
            stats_mode: if self.collect_stats {
                StatsMode::Basic as i32
            } else {
                StatsMode::None as i32
            },
            concurrent_result_sets: false,
            response_part_limit_bytes: 0,
            pool_id: String::new(),
            stats_period_ms: 0,
            schema_inclusion_mode: SchemaInclusionMode::Unspecified as i32,
            result_set_format: Format::Unspecified as i32,
            arrow_format_settings: None,
        })
    }
}

pub(crate) fn check_part(part: &ExecuteQueryResponsePart) -> RawResult<()> {
    check_status(part.status, &part.issues)
}

pub(crate) fn tx_id_from_part(part: &ExecuteQueryResponsePart) -> Option<String> {
    part.tx_meta
        .as_ref()
        .map(|m| m.id.clone())
        .filter(|id| !id.is_empty())
}

pub(crate) fn stats_from_part(part: &ExecuteQueryResponsePart) -> Option<Duration> {
    part.exec_stats
        .as_ref()
        .map(|stats| Duration::from_micros(stats.total_duration_us))
}

pub(crate) fn append_rows_from_part(
    columns: &mut Vec<RawColumn>,
    rows: &mut Vec<Vec<RawValue>>,
    truncated: &mut bool,
    part: ExecuteQueryResponsePart,
) -> RawResult<()> {
    let Some(proto_set) = part.result_set else {
        return Ok(());
    };
    let part_set = RawResultSet::try_from(proto_set)?;
    *truncated |= part_set.truncated;
    if !columns.is_empty()
        && !part_set.columns.is_empty()
        && !columns_compatible(columns, &part_set.columns)
    {
        return Err(crate::grpc_wrapper::raw_errors::RawError::custom(
            "column metadata mismatch between stream parts".to_string(),
        ));
    }
    if columns.is_empty() {
        *columns = part_set.columns;
    }
    rows.extend(part_set.rows);
    Ok(())
}

fn columns_compatible(existing: &[RawColumn], new_cols: &[RawColumn]) -> bool {
    existing.len() == new_cols.len()
        && existing
            .iter()
            .zip(new_cols.iter())
            .all(|(left, right)| left.name == right.name && left.column_type == right.column_type)
}