Skip to content

Commit

Permalink
snowflake-api: add non-critical error retries
Browse files Browse the repository at this point in the history
  • Loading branch information
andrusha committed Aug 18, 2023
1 parent d6489cb commit 792bfc7
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 9 deletions.
2 changes: 2 additions & 0 deletions snowflake-api/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }
Expand Down
2 changes: 1 addition & 1 deletion snowflake-api/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
29 changes: 21 additions & 8 deletions snowflake-api/src/connection.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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),

Expand Down Expand Up @@ -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>(
Expand Down

0 comments on commit 792bfc7

Please sign in to comment.