tushare_api/
client.rs

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/// HTTP client configuration for reqwest::Client
13#[derive(Debug, Clone)]
14pub struct HttpClientConfig {
15    /// Connection timeout duration
16    pub connect_timeout: Duration,
17    /// Request timeout duration
18    pub timeout: Duration,
19    /// Maximum idle connections per host
20    pub pool_max_idle_per_host: usize,
21    /// Pool idle timeout duration
22    pub pool_idle_timeout: Duration,
23    /// User agent string
24    pub user_agent: Option<String>,
25    /// Enable TCP_NODELAY to reduce latency
26    pub tcp_nodelay: bool,
27    /// TCP keep-alive duration
28    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,  // Increased for better performance
37            pool_idle_timeout: Duration::from_secs(90),  // Longer idle timeout
38            user_agent: Some("tushare-api-rust/1.0.0".to_string()),
39            tcp_nodelay: true,  // Reduce latency
40            tcp_keepalive: Some(Duration::from_secs(60)),  // Keep connections alive
41        }
42    }
43}
44
45impl HttpClientConfig {
46    /// Create a new HTTP client configuration with default values
47    pub fn new() -> Self {
48        Self::default()
49    }
50    
51    /// Set connection timeout
52    pub fn with_connect_timeout(mut self, timeout: Duration) -> Self {
53        self.connect_timeout = timeout;
54        self
55    }
56    
57    /// Set request timeout
58    pub fn with_timeout(mut self, timeout: Duration) -> Self {
59        self.timeout = timeout;
60        self
61    }
62    
63    /// Set maximum idle connections per host
64    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    /// Set pool idle timeout
70    pub fn with_pool_idle_timeout(mut self, timeout: Duration) -> Self {
71        self.pool_idle_timeout = timeout;
72        self
73    }
74    
75    /// Set user agent string
76    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    /// Enable or disable TCP_NODELAY
82    pub fn with_tcp_nodelay(mut self, enabled: bool) -> Self {
83        self.tcp_nodelay = enabled;
84        self
85    }
86    
87    /// Set TCP keep-alive duration
88    pub fn with_tcp_keepalive(mut self, duration: Option<Duration>) -> Self {
89        self.tcp_keepalive = duration;
90        self
91    }
92    
93    /// Build reqwest::Client with this configuration
94    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/// Internal request structure with token included
115#[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/// Tushare API client
136#[derive(Debug)]
137pub struct TushareClient {
138    token: String,
139    client: Client,
140    logger: Logger,
141}
142
143/// Tushare client builder
144#[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    /// Set HTTP client configuration
176    pub fn with_http_config(mut self, http_config: HttpClientConfig) -> Self {
177        self.http_config = http_config;
178        self
179    }
180    
181    /// Set maximum idle connections per host
182    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    /// Set pool idle timeout
188    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    /// Set log level
199    pub fn with_log_level(mut self, level: LogLevel) -> Self {
200        self.log_config.level = level;
201        self
202    }
203
204    /// Enable or disable request logging
205    pub fn log_requests(mut self, enabled: bool) -> Self {
206        self.log_config.log_requests = enabled;
207        self
208    }
209
210    /// Enable or disable response logging
211    pub fn log_responses(mut self, enabled: bool) -> Self {
212        self.log_config.log_responses = enabled;
213        self
214    }
215
216    /// Enable or disable sensitive data logging
217    pub fn log_sensitive_data(mut self, enabled: bool) -> Self {
218        self.log_config.log_sensitive_data = enabled;
219        self
220    }
221
222    /// Enable or disable performance metrics logging
223    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    /// Create client builder
244    pub fn builder() -> TushareClientBuilder {
245        TushareClientBuilder::new()
246    }
247
248    pub(crate) fn logger(&self) -> &Logger {
249        &self.logger
250    }
251
252
253
254    /// Create a new Tushare client with default timeout settings
255    /// 
256    /// # Arguments
257    /// 
258    /// * `token` - Tushare API Token
259    /// 
260    /// # Example
261    /// 
262    /// ```rust
263    /// use tushare_api::TushareClient;
264    /// 
265    /// let client = TushareClient::new("your_token_here");
266    /// ```
267    pub fn new(token: &str) -> Self {
268        Self::with_timeout(token, Duration::from_secs(10), Duration::from_secs(30))
269    }
270
271    /// Create a new Tushare client from TUSHARE_TOKEN environment variable with default timeout settings
272    /// 
273    /// # Errors
274    /// 
275    /// Returns `TushareError::InvalidToken` if TUSHARE_TOKEN environment variable does not exist or is empty
276    /// 
277    /// # Example
278    /// 
279    /// ```rust,no_run
280    /// use tushare_api::{TushareClient, TushareResult};
281    /// 
282    /// // Requires TUSHARE_TOKEN environment variable to be set
283    /// let client = TushareClient::from_env()?;
284    /// # Ok::<(), tushare_api::TushareError>(())
285    /// ```
286    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    /// Create a new Tushare client from TUSHARE_TOKEN environment variable with custom timeout settings
300    /// 
301    /// # Arguments
302    /// 
303    /// * `connect_timeout` - Connection timeout duration
304    /// * `timeout` - Request timeout duration
305    /// 
306    /// # Errors
307    /// 
308    /// Returns `TushareError::InvalidToken` if TUSHARE_TOKEN environment variable does not exist or is empty
309    /// 
310    /// # Example
311    /// 
312    /// ```rust,no_run
313    /// use tushare_api::{TushareClient, TushareResult};
314    /// use std::time::Duration;
315    /// 
316    /// // Requires TUSHARE_TOKEN environment variable to be set
317    /// let client = TushareClient::from_env_with_timeout(
318    ///     Duration::from_secs(5),  // Connection timeout 5 seconds
319    ///     Duration::from_secs(60)  // Request timeout 60 seconds
320    /// )?;
321    /// # Ok::<(), tushare_api::TushareError>(())
322    /// ```
323    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    /// Create a new Tushare client with custom timeout settings
337    /// 
338    /// # Arguments
339    /// 
340    /// * `token` - Tushare API Token
341    /// * `connect_timeout` - Connection timeout duration
342    /// * `timeout` - Request timeout duration
343    /// 
344    /// # Example
345    /// 
346    /// ```rust
347    /// use tushare_api::TushareClient;
348    /// use std::time::Duration;
349    /// 
350    /// let client = TushareClient::with_timeout(
351    ///     "your_token_here",
352    ///     Duration::from_secs(5),  // Connection timeout 5 seconds
353    ///     Duration::from_secs(60)  // Request timeout 60 seconds
354    /// );
355    /// ```
356    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    /// Call Tushare API with flexible string types support
372    /// 
373    /// # Arguments
374    /// 
375    /// * `request` - API request parameters, supports direct use of string literals
376    /// 
377    /// # Returns
378    /// 
379    /// Returns API response result
380    /// 
381    /// # Example
382    /// 
383    /// ```rust
384    /// use tushare_api::{TushareClient, TushareRequest, Api, params, fields, request};
385    /// 
386    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
387    ///     let client = TushareClient::new("your_token_here");
388    ///     
389    ///     // Now you can use string literals directly!
390    ///     let request = request!(Api::StockBasic, {
391    ///         "list_status" => "L"
392    ///     }, [
393    ///         "ts_code", "name"
394    ///     ]);
395    ///     
396    ///     let response = client.call_api(&request).await?;
397    ///     println!("Response: {:?}", response);
398    /// #   Ok(())
399    /// # }
400    /// ```
401    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        // Log API call start
433        self.logger.log_api_start(
434            &request_id,
435            &request.api_name.name(),
436            request.params.len(),
437            request.fields.len()
438        );
439        
440        // Log detailed request information (if enabled)
441        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        // Log success information and performance metrics
505        self.logger.log_api_success(&request_id, elapsed, tushare_response.data.clone().map(|data| data.items.len()).unwrap_or(0));
506        
507        // Log response details (if enabled)
508        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    /// 调用 Tushare API,并将响应的 `data.items` 解析为强类型的 [`TushareEntityList<T>`]。
518    ///
519    /// 这是 [`Self::call_api`] 的便捷封装:先执行请求,再把响应转换为实体列表。
520    ///
521    /// # Type Parameters
522    ///
523    /// - `T`: 单行数据对应的实体类型(需要实现 [`crate::traits::FromTushareData`])。
524    /// - `R`: 请求类型(需要实现 `TryInto<TushareRequest>`),通常可由参数自动推导。
525    ///
526    /// # Errors
527    ///
528    /// - 请求构造失败、网络/HTTP 错误、JSON/数据映射失败等都会以 [`TushareError`] 返回。
529    ///
530    /// # Example
531    ///
532    /// ```rust
533    /// # use tushare_api::{TushareClient, TushareRequest, TushareEntityList, Api, request, DeriveFromTushareData, params, fields};
534    /// # #[derive(Debug, Clone, DeriveFromTushareData)]
535    /// # struct Stock { ts_code: String }
536    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
537    /// let client = TushareClient::from_env()?;
538    /// let stocks: TushareEntityList<Stock> = client
539    ///     .call_api_as(request!(Api::StockBasic, {}, ["ts_code"]))
540    ///     .await?;
541    /// # Ok(()) }
542    /// ```
543    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 /// Generate a unique request ID for logging purposes
555pub(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        // let parmas = params!(
587        //     "list_status" => "L",
588        //     "limit" => "100"
589        // );
590        // let req = TushareRequest::new(Api::StockBasic, parmas, fields!("ts_code"));
591        // let response = client.call_api(req).await.unwrap();
592        // println!("resposne = {:?}", response);
593    }
594}