cronback_lib/clients/
scheduler_client.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use derive_more::{Deref, DerefMut};
5use proto::scheduler_proto::scheduler_client::SchedulerClient as GenSchedulerClient;
6use tonic::codegen::InterceptedService;
7
8use crate::config::MainConfig;
9use crate::grpc_client_provider::{GrpcClientType, ScopedGrpcClient};
10use crate::grpc_helpers::GrpcRequestInterceptor;
11use crate::prelude::ValidShardedId;
12use crate::types::{ProjectId, RequestId};
13
14type SchedulerClient = GenSchedulerClient<
15    InterceptedService<tonic::transport::Channel, GrpcRequestInterceptor>,
16>;
17
18#[derive(Deref, DerefMut)]
19pub struct ScopedSchedulerClient(ScopedGrpcClient<SchedulerClient>);
20
21#[async_trait]
22impl GrpcClientType for ScopedSchedulerClient {
23    type RawGrpcClient = SchedulerClient;
24
25    fn get_mut(&mut self) -> &mut ScopedGrpcClient<Self::RawGrpcClient> {
26        &mut self.0
27    }
28
29    fn address_map(config: &MainConfig) -> &HashMap<u64, String> {
30        &config.scheduler_cell_map
31    }
32
33    fn create_scoped_client(
34        project_id: ValidShardedId<ProjectId>,
35        request_id: RequestId,
36        channel: tonic::transport::Channel,
37        interceptor: GrpcRequestInterceptor,
38    ) -> Self {
39        let client = GenSchedulerClient::with_interceptor(channel, interceptor);
40
41        Self(ScopedGrpcClient::new(project_id, request_id, client))
42    }
43}