From 792bfc76d25bcd35ed93696afb30d3f6d3100149 Mon Sep 17 00:00:00 2001
From: Andrew Korzhuev <korzhuev@andrusha.me>
Date: Fri, 18 Aug 2023 17:26:40 +0200
Subject: [PATCH] snowflake-api: add non-critical error retries

---
 snowflake-api/Cargo.toml        |  2 ++
 snowflake-api/README.md         |  2 +-
 snowflake-api/src/connection.rs | 29 +++++++++++++++++++++--------
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/snowflake-api/Cargo.toml b/snowflake-api/Cargo.toml
index 0f8193a..f1dcc9b 100644
--- a/snowflake-api/Cargo.toml
+++ b/snowflake-api/Cargo.toml
@@ -15,6 +15,8 @@ license = "Apache-2.0"
 thiserror = "1"
 snowflake-jwt = "0.1.1"
 reqwest = { version = "0.11", features = ["json"] }
+reqwest-middleware = "0.2"
+reqwest-retry = "0.2"
 log = "0.4"
 serde_json = "1"
 serde = { version = "1", features = ["derive"] }
diff --git a/snowflake-api/README.md b/snowflake-api/README.md
index 6fe4613..36de047 100644
--- a/snowflake-api/README.md
+++ b/snowflake-api/README.md
@@ -4,7 +4,7 @@ Snowflake library for undocumented public API. If you want to query documented p
 
 ## Features
 
-Since it does a lot of I/O the library is async-only, examples use [tokio](https://tokio.rs/) as a runtime.
+Since it does a lot of I/O the library is async-only, and currently has hard dependency on [tokio](https://tokio.rs/) as a runtime due to use of [reqwest](https://github.com/seanmonstar/reqwest).
 
 - [x] Single statements
 - [ ] Multiple statements
diff --git a/snowflake-api/src/connection.rs b/snowflake-api/src/connection.rs
index 39ec913..efc647d 100644
--- a/snowflake-api/src/connection.rs
+++ b/snowflake-api/src/connection.rs
@@ -1,5 +1,8 @@
+use reqwest::header;
 use reqwest::header::{HeaderMap, HeaderValue};
-use reqwest::{header, Client, ClientBuilder};
+use reqwest_middleware::ClientWithMiddleware;
+use reqwest_retry::policies::ExponentialBackoff;
+use reqwest_retry::RetryTransientMiddleware;
 use thiserror::Error;
 use url::Url;
 use uuid::Uuid;
@@ -9,6 +12,9 @@ pub enum ConnectionError {
     #[error(transparent)]
     RequestError(#[from] reqwest::Error),
 
+    #[error(transparent)]
+    RequestMiddlewareError(#[from] reqwest_middleware::Error),
+
     #[error(transparent)]
     UrlParsing(#[from] url::ParseError),
 
@@ -55,24 +61,31 @@ impl QueryType {
 /// Minimal session will have at least 2 requests - login and query
 pub struct Connection {
     // no need for Arc as it's already inside the reqwest client
-    client: Client,
+    client: ClientWithMiddleware,
 }
 
 impl Connection {
     pub fn new() -> Result<Self, ConnectionError> {
+        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(3);
+
         // use builder to fail safely, unlike client new
-        let client = ClientBuilder::new()
+        let client = reqwest::ClientBuilder::new()
             .user_agent("Rust/0.0.1")
-            .referer(false)
-            // fixme: disable later
-            .connection_verbose(true)
-            .build()?;
+            .referer(false);
+
+        #[cfg(debug_assertions)]
+        let client = client.connection_verbose(true);
+
+        let client = client.build()?;
+
+        let client = reqwest_middleware::ClientBuilder::new(client)
+            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
+            .build();
 
         Ok(Connection { client })
     }
 
     /// Perform request of given query type with extra body or parameters
-    // todo: implement retry logic
     // todo: implement soft error handling
     // todo: is there better way to not repeat myself?
     pub async fn request<R: serde::de::DeserializeOwned>(