1use reqwest::Client;
2use std::time::{Duration, Instant};
3use std::collections::HashMap;
4use crate::error::{TushareError, TushareResult};
5use crate::types::{TushareRequest, TushareResponse, TushareEntityList};
6use crate::api::{Api, serialize_api_name};
7use crate::logging::{LogConfig, LogLevel, Logger};
8use serde::{Serialize};
9use serde_json;
10use std::time::{SystemTime, UNIX_EPOCH};
11
12#[derive(Debug, Clone)]
14pub struct HttpClientConfig {
15 pub connect_timeout: Duration,
17 pub timeout: Duration,
19 pub pool_max_idle_per_host: usize,
21 pub pool_idle_timeout: Duration,
23 pub user_agent: Option<String>,
25 pub tcp_nodelay: bool,
27 pub tcp_keepalive: Option<Duration>,
29}
30
31impl Default for HttpClientConfig {
32 fn default() -> Self {
33 Self {
34 connect_timeout: Duration::from_secs(10),
35 timeout: Duration::from_secs(30),
36 pool_max_idle_per_host: 20, pool_idle_timeout: Duration::from_secs(90), user_agent: Some("tushare-api-rust/1.0.0".to_string()),
39 tcp_nodelay: true, tcp_keepalive: Some(Duration::from_secs(60)), }
42 }
43}
44
45impl HttpClientConfig {
46 pub fn new() -> Self {
48 Self::default()
49 }
50
51 pub fn with_connect_timeout(mut self, timeout: Duration) -> Self {
53 self.connect_timeout = timeout;
54 self
55 }
56
57 pub fn with_timeout(mut self, timeout: Duration) -> Self {
59 self.timeout = timeout;
60 self
61 }
62
63 pub fn with_pool_max_idle_per_host(mut self, max_idle: usize) -> Self {
65 self.pool_max_idle_per_host = max_idle;
66 self
67 }
68
69 pub fn with_pool_idle_timeout(mut self, timeout: Duration) -> Self {
71 self.pool_idle_timeout = timeout;
72 self
73 }
74
75 pub fn with_user_agent<S: Into<String>>(mut self, user_agent: S) -> Self {
77 self.user_agent = Some(user_agent.into());
78 self
79 }
80
81 pub fn with_tcp_nodelay(mut self, enabled: bool) -> Self {
83 self.tcp_nodelay = enabled;
84 self
85 }
86
87 pub fn with_tcp_keepalive(mut self, duration: Option<Duration>) -> Self {
89 self.tcp_keepalive = duration;
90 self
91 }
92
93 pub(crate) fn build_client(&self) -> Result<Client, reqwest::Error> {
95 let mut builder = Client::builder()
96 .connect_timeout(self.connect_timeout)
97 .timeout(self.timeout)
98 .pool_max_idle_per_host(self.pool_max_idle_per_host)
99 .pool_idle_timeout(self.pool_idle_timeout)
100 .tcp_nodelay(self.tcp_nodelay);
101
102 if let Some(ref user_agent) = self.user_agent {
103 builder = builder.user_agent(user_agent);
104 }
105
106 if let Some(keepalive) = self.tcp_keepalive {
107 builder = builder.tcp_keepalive(keepalive);
108 }
109
110 builder.build()
111 }
112}
113
114#[derive(Debug)]
116struct ApiNameRef<'a>(&'a Api);
117
118impl<'a> Serialize for ApiNameRef<'a> {
119 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
120 where
121 S: serde::Serializer,
122 {
123 serialize_api_name(self.0, serializer)
124 }
125}
126
127#[derive(Debug, Serialize)]
128struct InternalTushareRequest<'a> {
129 api_name: ApiNameRef<'a>,
130 token: &'a str,
131 params: &'a HashMap<String, String>,
132 fields: &'a [String],
133}
134
135#[derive(Debug)]
137pub struct TushareClient {
138 token: String,
139 client: Client,
140 logger: Logger,
141}
142
143#[derive(Debug)]
145pub struct TushareClientBuilder {
146 token: Option<String>,
147 http_config: HttpClientConfig,
148 log_config: LogConfig,
149}
150
151impl TushareClientBuilder {
152 pub fn new() -> Self {
153 Self {
154 token: None,
155 http_config: HttpClientConfig::default(),
156 log_config: LogConfig::default(),
157 }
158 }
159
160 pub fn with_token(mut self, token: &str) -> Self {
161 self.token = Some(token.to_string());
162 self
163 }
164
165 pub fn with_connect_timeout(mut self, connect_timeout: Duration) -> Self {
166 self.http_config = self.http_config.with_connect_timeout(connect_timeout);
167 self
168 }
169
170 pub fn with_timeout(mut self, timeout: Duration) -> Self {
171 self.http_config = self.http_config.with_timeout(timeout);
172 self
173 }
174
175 pub fn with_http_config(mut self, http_config: HttpClientConfig) -> Self {
177 self.http_config = http_config;
178 self
179 }
180
181 pub fn with_pool_max_idle_per_host(mut self, max_idle: usize) -> Self {
183 self.http_config = self.http_config.with_pool_max_idle_per_host(max_idle);
184 self
185 }
186
187 pub fn with_pool_idle_timeout(mut self, timeout: Duration) -> Self {
189 self.http_config = self.http_config.with_pool_idle_timeout(timeout);
190 self
191 }
192
193 pub fn with_log_config(mut self, log_config: LogConfig) -> Self {
194 self.log_config = log_config;
195 self
196 }
197
198 pub fn with_log_level(mut self, level: LogLevel) -> Self {
200 self.log_config.level = level;
201 self
202 }
203
204 pub fn log_requests(mut self, enabled: bool) -> Self {
206 self.log_config.log_requests = enabled;
207 self
208 }
209
210 pub fn log_responses(mut self, enabled: bool) -> Self {
212 self.log_config.log_responses = enabled;
213 self
214 }
215
216 pub fn log_sensitive_data(mut self, enabled: bool) -> Self {
218 self.log_config.log_sensitive_data = enabled;
219 self
220 }
221
222 pub fn log_performance(mut self, enabled: bool) -> Self {
224 self.log_config.log_performance = enabled;
225 self
226 }
227
228 pub fn build(self) -> TushareResult<TushareClient> {
229 let token = self.token.ok_or(TushareError::InvalidToken)?;
230
231 let client = self.http_config.build_client()
232 .map_err(TushareError::HttpError)?;
233
234 Ok(TushareClient {
235 token,
236 client,
237 logger: Logger::new(self.log_config),
238 })
239 }
240}
241
242impl TushareClient {
243 pub fn builder() -> TushareClientBuilder {
245 TushareClientBuilder::new()
246 }
247
248 pub(crate) fn logger(&self) -> &Logger {
249 &self.logger
250 }
251
252
253
254 pub fn new(token: &str) -> Self {
268 Self::with_timeout(token, Duration::from_secs(10), Duration::from_secs(30))
269 }
270
271 pub fn from_env() -> TushareResult<Self> {
287 let token = std::env::var("TUSHARE_TOKEN")
288 .map_err(|_| TushareError::InvalidToken)?
289 .trim()
290 .to_string();
291
292 if token.is_empty() {
293 return Err(TushareError::InvalidToken);
294 }
295
296 Ok(Self::new(&token))
297 }
298
299 pub fn from_env_with_timeout(connect_timeout: Duration, timeout: Duration) -> TushareResult<Self> {
324 let token = std::env::var("TUSHARE_TOKEN")
325 .map_err(|_| TushareError::InvalidToken)?
326 .trim()
327 .to_string();
328
329 if token.is_empty() {
330 return Err(TushareError::InvalidToken);
331 }
332
333 Ok(Self::with_timeout(&token, connect_timeout, timeout))
334 }
335
336 pub fn with_timeout(token: &str, connect_timeout: Duration, timeout: Duration) -> Self {
357 let http_config = HttpClientConfig::new()
358 .with_connect_timeout(connect_timeout)
359 .with_timeout(timeout);
360
361 let client = http_config.build_client()
362 .expect("Failed to create HTTP client");
363
364 TushareClient {
365 token: token.to_string(),
366 client,
367 logger: Logger::new(LogConfig::default()),
368 }
369 }
370
371 pub async fn call_api<T>(&self, request: &T) -> TushareResult<TushareResponse>
402 where
403 for<'a> &'a T: TryInto<TushareRequest>,
404 for<'a> <&'a T as TryInto<TushareRequest>>::Error: Into<TushareError>,
405 {
406 let request = request
407 .try_into()
408 .map_err(Into::into)?;
409 let request_id = generate_request_id();
410 self.call_api_inner_with_request_id(&request_id, &request).await
411 }
412
413 pub(crate) async fn call_api_request(&self, request: &TushareRequest) -> TushareResult<TushareResponse> {
414 let request_id = generate_request_id();
415 self.call_api_inner_with_request_id(&request_id, request).await
416 }
417
418 pub(crate) async fn call_api_request_with_request_id(
419 &self,
420 request_id: &str,
421 request: &TushareRequest,
422 ) -> TushareResult<TushareResponse> {
423 self.call_api_inner_with_request_id(request_id, request).await
424 }
425
426 async fn call_api_inner_with_request_id(
427 &self,
428 request_id: &str,
429 request: &TushareRequest,
430 ) -> TushareResult<TushareResponse> {
431 let start_time = Instant::now();
432 self.logger.log_api_start(
434 &request_id,
435 &request.api_name.name(),
436 request.params.len(),
437 request.fields.len()
438 );
439
440 let token_preview_string = if self.logger.config().log_sensitive_data {
442 Some(format!("token: {}***", &self.token[..self.token.len().min(8)]))
443 } else {
444 None
445 };
446
447 self.logger.log_request_details(
448 &request_id,
449 &request.api_name.name(),
450 &format!("{:?}", request.params),
451 &format!("{:?}", request.fields),
452 token_preview_string.as_deref()
453 );
454
455 let internal_request = InternalTushareRequest {
456 api_name: ApiNameRef(&request.api_name),
457 token: &self.token,
458 params: &request.params,
459 fields: &request.fields,
460 };
461
462 self.logger.log_http_request(&request_id);
463
464 let response = self.client
465 .post("http://api.tushare.pro")
466 .json(&internal_request)
467 .send()
468 .await
469 .map_err(|e| {
470 let elapsed = start_time.elapsed();
471 self.logger.log_http_error(&request_id, elapsed, &e.to_string());
472 e
473 })?;
474
475 let status = response.status();
476 self.logger.log_http_response(&request_id, status.as_u16());
477
478 let response_text = response.text().await
479 .map_err(|e| {
480 let elapsed = start_time.elapsed();
481 self.logger.log_response_read_error(&request_id, elapsed, &e.to_string());
482 e
483 })?;
484 self.logger.log_raw_response(&request_id, &response_text);
485
486 let tushare_response: TushareResponse = serde_json::from_str(&response_text)
487 .map_err(|e| {
488 let elapsed = start_time.elapsed();
489 self.logger.log_json_parse_error(&request_id, elapsed, &e.to_string(), &response_text);
490 e
491 })?;
492
493 let elapsed = start_time.elapsed();
494
495 if tushare_response.code != 0 {
496 let message = format!("error code: {}, error msg: {}", tushare_response.code, tushare_response.msg.clone().unwrap_or_default());
497 self.logger.log_api_error(&request_id, elapsed, tushare_response.code, &message);
498 return Err(TushareError::ApiError {
499 code: tushare_response.code,
500 message
501 });
502 }
503
504 self.logger.log_api_success(&request_id, elapsed, tushare_response.data.clone().map(|data| data.items.len()).unwrap_or(0));
506
507 self.logger.log_response_details(
509 &request_id,
510 &tushare_response.request_id,
511 &format!("{:?}", tushare_response.data.as_ref().map(|d| &d.fields))
512 );
513
514 Ok(tushare_response)
515 }
516
517 pub async fn call_api_as<T, R>(&self, request: R) -> TushareResult<TushareEntityList<T>>
544 where
545 T: crate::traits::FromTushareData,
546 for<'a> &'a R: TryInto<TushareRequest>,
547 for<'a> <&'a R as TryInto<TushareRequest>>::Error: Into<TushareError>,
548 {
549 let response = self.call_api(&request).await?;
550 TushareEntityList::try_from(response).map_err(Into::into)
551 }
552 }
553
554 pub(crate) fn generate_request_id() -> String {
556 let timestamp = SystemTime::now()
557 .duration_since(UNIX_EPOCH)
558 .unwrap_or_default()
559 .as_nanos();
560 format!("req_{}", timestamp)
561}
562
563 mod tests {
564 use crate::{fields, params, Api, TushareClient, TushareRequest};
565
566 #[tokio::test]
567 async fn test() {
568 unsafe { std::env::set_var("TUSHARE_TOKEN", "xxxx"); }
569 let client = TushareClient::from_env().unwrap();
570 let response = client.call_api(&r#"
571 {
572 "api_name": "stock_basic",
573 "params": { "list_stauts": "L"},
574 "fields": [ "ts_code",
575 "symbol",
576 "name",
577 "area",
578 "industry",
579 "list_date",
580 "exchange",
581 "market"]
582 }
583 "#
584 ).await;
585 println!("resposne = {:?}", response);
586 }
594}