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
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

use ydb_grpc::ydb_proto::status_ids::StatusCode;

use crate::errors::{YdbError, YdbResult};
use crate::grpc_connection_manager::GrpcConnectionManager;
use crate::grpc_wrapper::raw_operation_service::client::RawOperationClient;
use crate::grpc_wrapper::raw_operation_service::types::RawListOperationsRequest;

use super::builders::{
    CancelOperationBuilder, ForgetOperationBuilder, GetOperationBuilder, ListOperationsBuilder,
    OperationCallOptions, raw_to_list_result, raw_to_operation_info, retry_operation_call,
};
use super::types::{ListOperationsRequest, ListOperationsResult, OperationInfo};
use crate::retry_budget::RetryControl;
use tracing::instrument;

#[derive(Clone)]
pub struct OperationClient {
    connection_manager: GrpcConnectionManager,
    retry_control: std::sync::Arc<RetryControl>,
}

impl OperationClient {
    pub(crate) fn new(
        connection_manager: GrpcConnectionManager,
        retry_control: std::sync::Arc<RetryControl>,
    ) -> Self {
        Self {
            connection_manager,
            retry_control,
        }
    }

    pub fn get_operation(&self, id: impl Into<String>) -> GetOperationBuilder<'_> {
        GetOperationBuilder {
            client: self,
            id: id.into(),
            opts: OperationCallOptions::default(),
        }
    }

    #[instrument(name = "ydb.OperationClient.GetOperation", skip_all, fields(db.system.name = "ydb", ydb.operation.id = %id), err)]
    pub(crate) async fn get_operation_call(
        &self,
        id: String,
        opts: OperationCallOptions,
    ) -> YdbResult<OperationInfo> {
        retry_operation_call(&self.retry_control, &opts, || async {
            let mut client = self.raw_client().await?;
            let op = client.get_operation(&id).await.map_err(YdbError::from)?;
            Ok(raw_to_operation_info(op))
        })
        .await
    }

    pub fn list_operations(&self, request: ListOperationsRequest) -> ListOperationsBuilder<'_> {
        ListOperationsBuilder {
            client: self,
            request,
            opts: OperationCallOptions::default(),
        }
    }

    #[instrument(name = "ydb.OperationClient.ListOperations", skip_all, fields(db.system.name = "ydb"), err)]
    pub(crate) async fn list_operations_call(
        &self,
        request: ListOperationsRequest,
        opts: OperationCallOptions,
    ) -> YdbResult<ListOperationsResult> {
        let raw_req = RawListOperationsRequest {
            kind: request.kind,
            page_size: request.page_size,
            page_token: request.page_token,
        };
        retry_operation_call(&self.retry_control, &opts, || async {
            let mut client = self.raw_client().await?;
            let result = client
                .list_operations(raw_req.clone())
                .await
                .map_err(YdbError::from)?;
            Ok(raw_to_list_result(result))
        })
        .await
    }

    /// Forgets a completed operation on the server.
    ///
    /// If the operation was already forgotten (e.g. a retry after a successful first attempt
    /// that lost the response), `NOT_FOUND` is treated as success.
    pub fn forget_operation(&self, id: impl Into<String>) -> ForgetOperationBuilder<'_> {
        ForgetOperationBuilder {
            client: self,
            id: id.into(),
            opts: OperationCallOptions::default(),
        }
    }

    #[instrument(name = "ydb.OperationClient.ForgetOperation", skip_all, fields(db.system.name = "ydb", ydb.operation.id = %id), err)]
    pub(crate) async fn forget_operation_call(
        &self,
        id: String,
        opts: OperationCallOptions,
    ) -> YdbResult<()> {
        retry_operation_call(&self.retry_control, &opts, || async {
            let mut client = self.raw_client().await?;
            match client.forget_operation(&id).await.map_err(YdbError::from) {
                Ok(()) => Ok(()),
                Err(YdbError::YdbStatusError(status))
                    if status.operation_status == StatusCode::NotFound as i32 =>
                {
                    Ok(())
                }
                Err(err) => Err(err),
            }
        })
        .await
    }

    pub fn cancel_operation(&self, id: impl Into<String>) -> CancelOperationBuilder<'_> {
        CancelOperationBuilder {
            client: self,
            id: id.into(),
            opts: OperationCallOptions::default(),
        }
    }

    #[instrument(name = "ydb.OperationClient.CancelOperation", skip_all, fields(db.system.name = "ydb", ydb.operation.id = %id), err)]
    pub(crate) async fn cancel_operation_call(
        &self,
        id: String,
        opts: OperationCallOptions,
    ) -> YdbResult<()> {
        retry_operation_call(&self.retry_control, &opts, || async {
            let mut client = self.raw_client().await?;
            client.cancel_operation(&id).await.map_err(YdbError::from)?;
            Ok(())
        })
        .await
    }

    async fn raw_client(&self) -> YdbResult<RawOperationClient> {
        self.connection_manager
            .get_auth_service(RawOperationClient::new)
            .await
    }
}