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

use std::sync::Arc;

use crate::connection_pool::{Connection, ConnectionPool, RacyRoundRobin, Simple};
use crate::grpc_wrapper::grpc_limits::WithGrpcMaxMessageSize;
use crate::grpc_wrapper::raw_services::{GrpcServiceForDiscovery, Service};
use crate::grpc_wrapper::runtime_interceptors::{InterceptedChannel, MultiInterceptor};
use crate::load_balancer::{LoadBalancer, SharedLoadBalancer};
use crate::{GrpcOptions, YdbResult};
use derivative::Derivative;
use http::Uri;
use tracing::instrument;

pub(crate) type GrpcConnectionManager = GrpcConnectionManagerGeneric<SharedLoadBalancer, Simple>;
pub(crate) type DiscoveryConnectionManager =
    GrpcConnectionManagerGeneric<NoBalancer, RacyRoundRobin>;

#[derive(Debug, Clone, Copy)]
pub(crate) struct NoBalancer;

#[derive(Derivative)]
#[derivative(Clone(bound = "BalancerT: Clone"), Debug)]
pub(crate) struct GrpcConnectionManagerGeneric<BalancerT, ConnectionT: Connection> {
    balancer: BalancerT,
    connections_pool: Arc<ConnectionPool<ConnectionT>>,
    #[derivative(Debug = "ignore")]
    interceptor: MultiInterceptor,
    database: String,
    opts: GrpcOptions,
}

impl<BalancerT, ConnectionT: Connection> GrpcConnectionManagerGeneric<BalancerT, ConnectionT> {
    pub(crate) fn new(
        balancer: BalancerT,
        database: String,
        interceptor: MultiInterceptor,
        opts: GrpcOptions,
    ) -> Self {
        let cp = ConnectionPool::new(opts.clone());

        Self {
            balancer,
            connections_pool: cp.into(),
            interceptor,
            database,
            opts,
        }
    }

    #[instrument(name = "ydb.ConnectionManager.GetAuthService", skip_all, fields(db.system.name = "ydb", ydb.service.name = ?T::get_grpc_discovery_service()))]
    pub(crate) async fn get_auth_service<
        T: GrpcServiceForDiscovery + WithGrpcMaxMessageSize,
        F: FnOnce(InterceptedChannel) -> T,
    >(
        &self,
        new: F,
    ) -> YdbResult<T>
    where
        BalancerT: LoadBalancer,
    {
        let uri = self.balancer.endpoint(T::get_grpc_discovery_service())?;
        self.get_auth_service_to_node(new, &uri).await
    }

    #[instrument(name = "ydb.ConnectionManager.GetAuthServiceToNode", skip_all)]
    pub(crate) async fn get_auth_service_to_node<
        T: GrpcServiceForDiscovery + WithGrpcMaxMessageSize,
        F: FnOnce(InterceptedChannel) -> T,
    >(
        &self,
        new: F,
        uri: &Uri,
    ) -> YdbResult<T> {
        let channel = self.connections_pool.connection(uri).await?;

        let intercepted_channel = InterceptedChannel::new(channel, self.interceptor.clone());
        Ok(new(intercepted_channel).with_grpc_max_message_size(self.opts.max_message_size))
    }

    #[instrument(name = "ydb.ConnectionManager.GetEndpoint", skip_all, fields(ydb.service.name = ?service))]
    pub(crate) fn endpoint(&self, service: Service) -> YdbResult<Uri>
    where
        BalancerT: LoadBalancer,
    {
        self.balancer.endpoint(service)
    }

    pub(crate) fn database(&self) -> &String {
        &self.database
    }
}