rigg_client/
client.rs

1//! Azure Search REST API client
2
3use std::time::Duration;
4
5use reqwest::{Client, Method, StatusCode};
6use serde_json::Value;
7use tracing::{debug, instrument, warn};
8
9use rigg_core::config::SearchServiceConfig;
10use rigg_core::resources::ResourceKind;
11
12use crate::auth::{AuthProvider, get_auth_provider};
13use crate::error::ClientError;
14
15/// Maximum number of retry attempts for retryable errors
16const MAX_RETRIES: u32 = 3;
17
18/// Initial backoff delay in seconds
19const INITIAL_BACKOFF_SECS: u64 = 1;
20
21/// Calculate the backoff duration for a given retry attempt.
22///
23/// For `RateLimited` errors with a `retry_after` value, that value is used directly.
24/// For other retryable errors, exponential backoff is applied: 1s, 2s, 4s, etc.
25fn retry_delay(error: &ClientError, attempt: u32) -> Duration {
26    match error {
27        ClientError::RateLimited { retry_after } => Duration::from_secs(*retry_after),
28        _ => Duration::from_secs(INITIAL_BACKOFF_SECS * 2u64.pow(attempt)),
29    }
30}
31
32/// Azure Search API client
33pub struct AzureSearchClient {
34    http: Client,
35    auth: Box<dyn AuthProvider>,
36    base_url: String,
37    api_version: String,
38    preview_api_version: String,
39}
40
41impl AzureSearchClient {
42    /// Create client from a workspace search connection.
43    pub fn from_connection(
44        conn: &rigg_core::workspace::SearchConnection,
45    ) -> Result<Self, ClientError> {
46        let auth = get_auth_provider()?;
47        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
48        Ok(Self {
49            http,
50            auth,
51            base_url: conn
52                .endpoint
53                .clone()
54                .map(|e| e.trim_end_matches('/').to_string())
55                .unwrap_or_else(|| format!("https://{}.search.windows.net", conn.service)),
56            api_version: conn
57                .api_version
58                .clone()
59                .unwrap_or_else(|| rigg_core::registry::SEARCH_STABLE_API_VERSION.to_string()),
60            preview_api_version: conn
61                .preview_api_version
62                .clone()
63                .unwrap_or_else(|| rigg_core::registry::SEARCH_PREVIEW_API_VERSION.to_string()),
64        })
65    }
66
67    /// Create client from a specific search service config (legacy).
68    pub fn from_service_config(service: &SearchServiceConfig) -> Result<Self, ClientError> {
69        let auth = get_auth_provider()?;
70        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
71
72        Ok(Self {
73            http,
74            auth,
75            base_url: service.service_url(),
76            api_version: rigg_core::registry::SEARCH_STABLE_API_VERSION.to_string(),
77            preview_api_version: service.preview_api_version.clone(),
78        })
79    }
80
81    /// Create with a custom auth provider and explicit versions (tests).
82    pub fn with_auth(
83        base_url: String,
84        api_version: String,
85        preview_api_version: String,
86        auth: Box<dyn AuthProvider>,
87    ) -> Result<Self, ClientError> {
88        let http = Client::builder()
89            .timeout(std::time::Duration::from_secs(30))
90            .build()?;
91
92        Ok(Self {
93            http,
94            auth,
95            base_url,
96            api_version,
97            preview_api_version,
98        })
99    }
100
101    /// Get the API version to use for a resource kind, per the registry
102    /// channel: stable kinds use the stable api-version, preview-gated kinds
103    /// the preview one. (Pre-0.18 the client always used preview; GA of
104    /// agentic retrieval in 2026-04-01 makes stable the default.)
105    fn api_version_for(&self, kind: ResourceKind) -> &str {
106        match rigg_core::registry::meta(kind).channel {
107            rigg_core::registry::Channel::Stable => &self.api_version,
108            rigg_core::registry::Channel::Preview => &self.preview_api_version,
109        }
110    }
111
112    /// Build URL for a resource collection
113    fn collection_url(&self, kind: ResourceKind) -> String {
114        format!(
115            "{}/{}?api-version={}",
116            self.base_url,
117            kind.api_path(),
118            self.api_version_for(kind)
119        )
120    }
121
122    /// Build URL for a specific resource
123    fn resource_url(&self, kind: ResourceKind, name: &str) -> String {
124        format!(
125            "{}/{}/{}?api-version={}",
126            self.base_url,
127            kind.api_path(),
128            urlencoding::encode(name),
129            self.api_version_for(kind)
130        )
131    }
132
133    /// Execute an HTTP request
134    async fn request(
135        &self,
136        method: Method,
137        url: &str,
138        body: Option<&Value>,
139    ) -> Result<Option<Value>, ClientError> {
140        let token = self.auth.get_token()?;
141
142        let mut request = self
143            .http
144            .request(method.clone(), url)
145            .header("Authorization", format!("Bearer {}", token))
146            .header("Content-Type", "application/json");
147
148        if let Some(json) = body {
149            request = request.json(json);
150        }
151
152        debug!("Request: {} {}", method, url);
153        let response = request.send().await?;
154        let status = response.status();
155
156        if status == StatusCode::NO_CONTENT {
157            return Ok(None);
158        }
159
160        let body = response.text().await?;
161
162        if status.is_success() {
163            if body.is_empty() {
164                Ok(None)
165            } else {
166                let value: Value = serde_json::from_str(&body)?;
167                Ok(Some(value))
168            }
169        } else {
170            match status {
171                StatusCode::NOT_FOUND => Err(ClientError::NotFound {
172                    kind: "resource".to_string(),
173                    name: url.to_string(),
174                }),
175                StatusCode::CONFLICT => Err(ClientError::AlreadyExists {
176                    kind: "resource".to_string(),
177                    name: url.to_string(),
178                }),
179                StatusCode::TOO_MANY_REQUESTS => {
180                    let retry_after = 60; // Default retry time
181                    Err(ClientError::RateLimited { retry_after })
182                }
183                StatusCode::SERVICE_UNAVAILABLE => Err(ClientError::ServiceUnavailable(body)),
184                _ => Err(ClientError::from_response_with_url(
185                    status.as_u16(),
186                    &body,
187                    Some(url),
188                )),
189            }
190        }
191    }
192
193    /// Execute an HTTP request with retry logic for transient errors.
194    ///
195    /// Retries up to [`MAX_RETRIES`] times for retryable errors (429 and 503).
196    /// Uses exponential backoff (1s, 2s, 4s) for 503 errors and respects the
197    /// `retry_after` value for 429 rate-limiting errors.
198    async fn request_with_retry(
199        &self,
200        method: Method,
201        url: &str,
202        body: Option<&Value>,
203    ) -> Result<Option<Value>, ClientError> {
204        let mut attempt = 0u32;
205        loop {
206            match self.request(method.clone(), url, body).await {
207                Ok(value) => return Ok(value),
208                Err(err) if err.is_retryable() && attempt < MAX_RETRIES => {
209                    let delay = retry_delay(&err, attempt);
210                    warn!(
211                        "Request {} {} failed (attempt {}/{}): {}. Retrying in {:?}",
212                        method,
213                        url,
214                        attempt + 1,
215                        MAX_RETRIES + 1,
216                        err,
217                        delay,
218                    );
219                    tokio::time::sleep(delay).await;
220                    attempt += 1;
221                }
222                Err(err) => return Err(err),
223            }
224        }
225    }
226
227    /// List all resources of a given kind
228    #[instrument(skip(self))]
229    pub async fn list(&self, kind: ResourceKind) -> Result<Vec<Value>, ClientError> {
230        let url = self.collection_url(kind);
231        let response = self.request_with_retry(Method::GET, &url, None).await?;
232
233        match response {
234            Some(value) => {
235                // Azure returns { "value": [...] }
236                let items = value
237                    .get("value")
238                    .and_then(|v| v.as_array())
239                    .cloned()
240                    .unwrap_or_default();
241                Ok(items)
242            }
243            None => Ok(Vec::new()),
244        }
245    }
246
247    /// Get a specific resource
248    #[instrument(skip(self))]
249    pub async fn get(&self, kind: ResourceKind, name: &str) -> Result<Value, ClientError> {
250        let url = self.resource_url(kind, name);
251        let response = self.request_with_retry(Method::GET, &url, None).await?;
252
253        response.ok_or_else(|| ClientError::NotFound {
254            kind: kind.display_name().to_string(),
255            name: name.to_string(),
256        })
257    }
258
259    /// Create or update a resource
260    ///
261    /// Returns the response body if the API returns one. Some APIs (especially
262    /// preview endpoints like Knowledge Sources) return 204 No Content on
263    /// successful update, which yields `Ok(None)`.
264    #[instrument(skip(self, definition))]
265    pub async fn create_or_update(
266        &self,
267        kind: ResourceKind,
268        name: &str,
269        definition: &Value,
270    ) -> Result<Option<Value>, ClientError> {
271        let url = self.resource_url(kind, name);
272        self.request_with_retry(Method::PUT, &url, Some(definition))
273            .await
274    }
275
276    /// Delete a resource
277    #[instrument(skip(self))]
278    pub async fn delete(&self, kind: ResourceKind, name: &str) -> Result<(), ClientError> {
279        let url = self.resource_url(kind, name);
280        self.request_with_retry(Method::DELETE, &url, None).await?;
281        Ok(())
282    }
283
284    /// Check if a resource exists
285    pub async fn exists(&self, kind: ResourceKind, name: &str) -> Result<bool, ClientError> {
286        match self.get(kind, name).await {
287            Ok(_) => Ok(true),
288            Err(ClientError::NotFound { .. }) => Ok(false),
289            Err(e) => Err(e),
290        }
291    }
292
293    /// Get the authentication method being used
294    pub fn auth_method(&self) -> &'static str {
295        self.auth.method_name()
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302    use crate::auth::{AuthError, AuthProvider};
303
304    struct FakeAuth;
305    impl AuthProvider for FakeAuth {
306        fn get_token(&self) -> Result<String, AuthError> {
307            Ok("fake-token".to_string())
308        }
309        fn method_name(&self) -> &'static str {
310            "Fake"
311        }
312    }
313
314    fn make_client() -> AzureSearchClient {
315        AzureSearchClient::with_auth(
316            "https://test-svc.search.windows.net".to_string(),
317            "2026-04-01".to_string(),
318            "2026-05-01-preview".to_string(),
319            Box::new(FakeAuth),
320        )
321        .unwrap()
322    }
323
324    #[test]
325    fn stable_kinds_use_stable_version() {
326        let client = make_client();
327        let url = client.collection_url(ResourceKind::Index);
328        assert_eq!(
329            url,
330            "https://test-svc.search.windows.net/indexes?api-version=2026-04-01"
331        );
332        // agentic retrieval is GA in 2026-04-01
333        let url = client.collection_url(ResourceKind::KnowledgeBase);
334        assert_eq!(
335            url,
336            "https://test-svc.search.windows.net/knowledgeBases?api-version=2026-04-01"
337        );
338        let url = client.resource_url(ResourceKind::KnowledgeSource, "ks");
339        assert_eq!(
340            url,
341            "https://test-svc.search.windows.net/knowledgeSources/ks?api-version=2026-04-01"
342        );
343    }
344
345    #[test]
346    fn resource_url_percent_encodes_name() {
347        let client = make_client();
348        let url = client.resource_url(ResourceKind::Index, "my index");
349        assert!(url.contains("/indexes/my%20index?"));
350    }
351
352    #[test]
353    fn search_kinds_route_via_registry_channel() {
354        let client = make_client();
355        for kind in ResourceKind::search_kinds() {
356            let url = client.collection_url(kind);
357            let expected = match rigg_core::registry::meta(kind).channel {
358                rigg_core::registry::Channel::Stable => "2026-04-01",
359                rigg_core::registry::Channel::Preview => "2026-05-01-preview",
360            };
361            assert!(
362                url.contains(expected),
363                "{kind:?} should use {expected}, got: {url}"
364            );
365        }
366    }
367
368    #[test]
369    fn test_retry_delay_exponential_backoff_attempt_0() {
370        let err = ClientError::ServiceUnavailable("down".to_string());
371        let delay = retry_delay(&err, 0);
372        assert_eq!(delay, Duration::from_secs(1));
373    }
374
375    #[test]
376    fn test_retry_delay_exponential_backoff_attempt_1() {
377        let err = ClientError::ServiceUnavailable("down".to_string());
378        let delay = retry_delay(&err, 1);
379        assert_eq!(delay, Duration::from_secs(2));
380    }
381
382    #[test]
383    fn test_retry_delay_exponential_backoff_attempt_2() {
384        let err = ClientError::ServiceUnavailable("down".to_string());
385        let delay = retry_delay(&err, 2);
386        assert_eq!(delay, Duration::from_secs(4));
387    }
388
389    #[test]
390    fn test_retry_delay_rate_limited_uses_retry_after() {
391        let err = ClientError::RateLimited { retry_after: 30 };
392        // retry_after should be used regardless of attempt number
393        assert_eq!(retry_delay(&err, 0), Duration::from_secs(30));
394        assert_eq!(retry_delay(&err, 1), Duration::from_secs(30));
395        assert_eq!(retry_delay(&err, 2), Duration::from_secs(30));
396    }
397
398    #[test]
399    fn test_retry_delay_rate_limited_default_retry_after() {
400        let err = ClientError::RateLimited { retry_after: 60 };
401        let delay = retry_delay(&err, 0);
402        assert_eq!(delay, Duration::from_secs(60));
403    }
404
405    #[test]
406    fn test_retry_constants() {
407        assert_eq!(MAX_RETRIES, 3);
408        assert_eq!(INITIAL_BACKOFF_SECS, 1);
409    }
410
411    #[test]
412    fn test_retry_delay_backoff_sequence() {
413        let err = ClientError::ServiceUnavailable("temporarily unavailable".to_string());
414        let delays: Vec<Duration> = (0..MAX_RETRIES).map(|i| retry_delay(&err, i)).collect();
415        assert_eq!(
416            delays,
417            vec![
418                Duration::from_secs(1),
419                Duration::from_secs(2),
420                Duration::from_secs(4),
421            ]
422        );
423    }
424
425    #[test]
426    fn test_non_retryable_error_still_computes_delay() {
427        // retry_delay computes a delay regardless; the caller decides whether to retry.
428        // This verifies the function doesn't panic on non-retryable errors.
429        let err = ClientError::Api {
430            status: 400,
431            message: "bad request".to_string(),
432        };
433        let delay = retry_delay(&err, 0);
434        assert_eq!(delay, Duration::from_secs(1));
435    }
436}
rigg_client/client.rs

rigg_client/
client.rs
