tests_fuzz/utils/
kafka_wal_http.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use async_trait::async_trait;
16use serde::{Deserialize, Serialize};
17use snafu::ResultExt;
18
19use crate::error::{self, Result};
20
21#[derive(Debug, Clone, Deserialize, Serialize)]
22pub struct TopicOffset {
23    pub topic: String,
24    pub partition: i32,
25    pub offset: i64,
26}
27
28#[derive(Debug, Serialize, Deserialize)]
29pub struct RecordOffsetsRequest {
30    pub broker_endpoints: Vec<String>,
31    pub topic_prefix: String,
32    pub num_topics: usize,
33    #[serde(default)]
34    pub partition: i32,
35}
36
37#[derive(Debug, Serialize, Deserialize)]
38pub struct DeleteRecordsRequest {
39    pub broker_endpoints: Vec<String>,
40    pub offsets: Vec<TopicOffset>,
41    pub timeout_ms: Option<i32>,
42}
43
44#[derive(Debug, Serialize, Deserialize)]
45pub struct RecordOffsetsResponse {
46    pub offsets: Vec<TopicOffset>,
47}
48
49#[derive(Debug, Serialize, Deserialize)]
50pub struct DeleteRecordsResponse {
51    pub deleted: Vec<TopicOffset>,
52}
53
54#[async_trait]
55pub trait WalAdmin: Send + Sync {
56    async fn record_topic_latest_offsets(&self) -> Result<Vec<TopicOffset>>;
57
58    async fn delete_records_to_offsets(&self, offsets: &[TopicOffset]) -> Result<()>;
59}
60
61pub struct HttpWalAdmin {
62    client: reqwest::Client,
63    base_url: String,
64    broker_endpoints: Vec<String>,
65    topic_prefix: String,
66    num_topics: usize,
67    partition: i32,
68    delete_records_timeout_ms: i32,
69}
70
71pub struct HttpWalAdminConfig {
72    pub base_url: String,
73    pub broker_endpoints: Vec<String>,
74    pub topic_prefix: String,
75    pub num_topics: usize,
76    pub partition: i32,
77    pub delete_records_timeout_ms: i32,
78}
79
80impl HttpWalAdmin {
81    pub fn new(config: HttpWalAdminConfig) -> Self {
82        Self {
83            client: reqwest::Client::new(),
84            base_url: config.base_url,
85            broker_endpoints: config.broker_endpoints,
86            topic_prefix: config.topic_prefix,
87            num_topics: config.num_topics,
88            partition: config.partition,
89            delete_records_timeout_ms: config.delete_records_timeout_ms,
90        }
91    }
92
93    fn endpoint(&self, path: &str) -> String {
94        format!("{}{}", self.base_url.trim_end_matches('/'), path)
95    }
96}
97
98#[async_trait]
99impl WalAdmin for HttpWalAdmin {
100    async fn record_topic_latest_offsets(&self) -> Result<Vec<TopicOffset>> {
101        let request = RecordOffsetsRequest {
102            broker_endpoints: self.broker_endpoints.clone(),
103            topic_prefix: self.topic_prefix.clone(),
104            num_topics: self.num_topics,
105            partition: self.partition,
106        };
107        let response = self
108            .client
109            .post(self.endpoint("/record-offsets"))
110            .json(&request)
111            .send()
112            .await
113            .context(error::KafkaWalHelperRequestSnafu {
114                operation: "record-offsets".to_string(),
115            })?;
116        ensure_success(response, "record-offsets")
117            .await?
118            .json::<RecordOffsetsResponse>()
119            .await
120            .context(error::KafkaWalHelperRequestSnafu {
121                operation: "record-offsets".to_string(),
122            })
123            .map(|response| response.offsets)
124    }
125
126    async fn delete_records_to_offsets(&self, offsets: &[TopicOffset]) -> Result<()> {
127        let request = DeleteRecordsRequest {
128            broker_endpoints: self.broker_endpoints.clone(),
129            offsets: offsets.to_vec(),
130            timeout_ms: Some(self.delete_records_timeout_ms),
131        };
132        let response = self
133            .client
134            .post(self.endpoint("/delete-records"))
135            .json(&request)
136            .send()
137            .await
138            .context(error::KafkaWalHelperRequestSnafu {
139                operation: "delete-records".to_string(),
140            })?;
141        ensure_success(response, "delete-records")
142            .await?
143            .json::<DeleteRecordsResponse>()
144            .await
145            .context(error::KafkaWalHelperRequestSnafu {
146                operation: "delete-records".to_string(),
147            })
148            .map(|_| ())
149    }
150}
151
152async fn ensure_success(response: reqwest::Response, operation: &str) -> Result<reqwest::Response> {
153    let status = response.status();
154    if status.is_success() {
155        return Ok(response);
156    }
157
158    let body = response.text().await.unwrap_or_default();
159    error::KafkaWalHelperStatusSnafu {
160        operation: operation.to_string(),
161        status,
162        body,
163    }
164    .fail()
165}
tests_fuzz/utils/kafka_wal_http.rs

tests_fuzz/utils/
kafka_wal_http.rs
