feat(python): support customize credential provider for write_dataset…

… api in PyLance
lancedb · Dec 23, 2024 · 2653d0d · 2653d0d
1 parent b3219e8
commit 2653d0d
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/python/src/object_store.rs b/python/src/object_store.rs
@@ -53,7 +53,7 @@ impl UrlBasedCredentialProvider {
         let state = self.state.lock().await;
         if let Some(expiration) = state.expired_time {
             if SystemTime::now() < expiration - Duration::from_secs(600) {
-                Some(AwsCredential {
+                return Some(AwsCredential {
                     key_id: state.access_key_id.clone().unwrap_or_default(),
                     secret_key: state.secret_access_key.clone().unwrap_or_default(),
                     token: state.session_token.clone(),
@@ -84,14 +84,11 @@ impl CredentialProvider for UrlBasedCredentialProvider {
             return Ok(Arc::from(credentials));
         }
 
-        let response = match self.client.get(&self.url).send().await {
-            Ok(resp) => resp,
-            Err(_) => {
-                return Err(object_store::Error::Generic {
-                    store: "Request credential error.",
-                    source: Box::from(""),
-                })
-            }
+        let Ok(response) = self.client.get(&self.url).send().await else {
+            return Err(object_store::Error::Generic {
+                store: "Request credential error.",
+                source: Box::from(""),
+            });
         };
 
         let credential_state: CredentialState = match response.json().await {