Skip to content

Commit 68677da

Browse files
committed
feat: added auth headers to runtime status request
1 parent 625a7eb commit 68677da

File tree

2 files changed

+24
-11
lines changed

2 files changed

+24
-11
lines changed

crates/base/src/client/mod.rs

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
11
use std::time::{SystemTime, UNIX_EPOCH};
22

3+
use code0_flow::flow_service::auth::get_authorization_metadata;
34
use tokio::time::sleep;
4-
use tonic::transport::{Channel, Endpoint};
5+
use tonic::{
6+
Extensions, Request,
7+
transport::{Channel, Endpoint},
8+
};
59
use tucana::{
610
aquila::{
711
RuntimeStatusUpdateRequest, runtime_status_service_client::RuntimeStatusServiceClient,
@@ -15,6 +19,7 @@ pub struct DracoRuntimeStatusService {
1519
identifier: String,
1620
features: Vec<RuntimeFeature>,
1721
configs: Vec<AdapterConfiguration>,
22+
aquila_token: String,
1823
}
1924

2025
const MAX_BACKOFF: u64 = 2000 * 60;
@@ -70,22 +75,25 @@ impl DracoRuntimeStatusService {
7075
identifier: String,
7176
features: Vec<RuntimeFeature>,
7277
configs: Vec<AdapterConfiguration>,
78+
aquila_token: String
7379
) -> Self {
7480
let channel = create_channel_with_retry("Aquila", aquila_url).await;
75-
Self::new(channel, identifier, features, configs)
81+
Self::new(channel, identifier, features, configs, aquila_token)
7682
}
7783

7884
pub fn new(
7985
channel: Channel,
8086
identifier: String,
8187
features: Vec<RuntimeFeature>,
8288
configs: Vec<AdapterConfiguration>,
89+
aquila_token: String
8390
) -> Self {
8491
DracoRuntimeStatusService {
8592
channel,
8693
identifier,
8794
features,
8895
configs,
96+
aquila_token
8997
}
9098
}
9199

@@ -105,15 +113,19 @@ impl DracoRuntimeStatusService {
105113
}
106114
};
107115

108-
let request = RuntimeStatusUpdateRequest {
109-
status: Some(Status::AdapterRuntimeStatus(AdapterRuntimeStatus {
110-
status: status.into(),
111-
timestamp: timestamp as i64,
112-
identifier: self.identifier.clone(),
113-
features: self.features.clone(),
114-
configurations: self.configs.clone(),
115-
})),
116-
};
116+
let request = Request::from_parts(
117+
get_authorization_metadata(&self.aquila_token),
118+
Extensions::new(),
119+
RuntimeStatusUpdateRequest {
120+
status: Some(Status::AdapterRuntimeStatus(AdapterRuntimeStatus {
121+
status: status.into(),
122+
timestamp: timestamp as i64,
123+
identifier: self.identifier.clone(),
124+
features: self.features.clone(),
125+
configurations: self.configs.clone(),
126+
})),
127+
},
128+
);
117129

118130
match client.update(request).await {
119131
Ok(response) => {

crates/base/src/runner.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ impl<C: LoadConfig> ServerRunner<C> {
7373
config.draco_variant.clone(),
7474
runtime_feature,
7575
runtime_config,
76+
config.aquila_token.clone(),
7677
)
7778
.await,
7879
);

0 commit comments

Comments
 (0)