Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add dynamic URL provider #2426

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Add dynamic URL provider #2426

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b2ede6e
Add dynamic URL provider
ch-sc Jun 8, 2020
da1ed71
Use correct test connection
ch-sc Jun 8, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 47 additions & 1 deletion diesel/src/r2d2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use crate::deserialize::{Queryable, QueryableByName};
use crate::prelude::*;
use crate::query_builder::{AsQuery, QueryFragment, QueryId};
use crate::sql_types::HasSqlType;
use std::sync::Arc;

/// An r2d2 connection manager for use with Diesel.
///
Expand All @@ -31,6 +32,7 @@ use crate::sql_types::HasSqlType;
#[derive(Debug, Clone)]
pub struct ConnectionManager<T> {
database_url: String,
url_provider: Option<Arc<dyn UrlProvider>>,
_marker: PhantomData<T>,
}

Expand All @@ -42,6 +44,17 @@ impl<T> ConnectionManager<T> {
pub fn new<S: Into<String>>(database_url: S) -> Self {
ConnectionManager {
database_url: database_url.into(),
url_provider: None,
_marker: PhantomData,
}
}

/// Returns a new connection manager,
/// which establishes connections with the given UrlProvider
pub fn new_with_url_provider(url_provider: Arc<dyn UrlProvider>) -> Self {
ConnectionManager {
database_url: "".into(),
url_provider: Some(url_provider),
_marker: PhantomData,
}
}
Expand Down Expand Up @@ -103,7 +116,11 @@ where
type Error = Error;

fn connect(&self) -> Result<T, Error> {
T::establish(&self.database_url).map_err(Error::ConnectionError)
let db_url = match &self.url_provider {
Some(url_provider) => url_provider.provide_url(),
_ => self.database_url.to_owned(),
};
T::establish(&db_url).map_err(Error::ConnectionError)
}

fn is_valid(&self, conn: &mut T) -> Result<(), Error> {
Expand Down Expand Up @@ -175,6 +192,13 @@ where
}
}

/// Trait to provide the user the option to change parameters of the URL at runtime.
/// E.g. to implement password rotation
pub trait UrlProvider: Send + Sync + fmt::Debug + 'static {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why just do not use Fn/FnOnce/FnMut trait?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For better customizability of the provider.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For example a provider such as this one would enable password rotation.

#[derive(Clone, Debug)]
struct CustomUrlProvider {
    password: Arc<RwLock<String>>,
    host: String,
    port: u32,
    user: String,
    database: String,
    region: Region
}

impl CustomUrlProvider {

    // ...

    fn rotate_password(&self) -> Result<()> {
        let new_token = crate::iam_authentication::generate_db_auth_token(
            &self.region,
            &self.host,
            &self.port,
            &self.user,
            &WebIdentityProvider::from_k8s_env()
        );

        if let Ok(mut pw) = self.password.write() {
            *pw = crate::iam_authentication::urlencode(&new_token?);
            Ok(())
        } else {
            Err(CustomError::DbPasswordRotationError())
        }
    }
}

impl UrlProvider for CustomUrlProvider {
    fn provide_url(&self) -> String {
        loop {
            if let Ok(pw) = self.password.read() {
                break format!(
                    "postgres://{}:{}@{}:{}/{}",
                    self.user, *pw, self.host, self.port, self.database
                );
            } else {
                warn!("Couldn't acquire lock to build DB password")
            }
        }
    }
}

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it should be easy to just provide a default implementation all types that implement Fn() somehow.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, it's possible to use a closure instead. Nonetheless, to me it seems more handy to use a new trait. A custom provider can have a state and easily be dealing with connection updates or different test scenarios.

On the contrary, it's really just about making the db connection more dynamic and a Fn() would suffice.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be clear: I'm not talking about removing UrlProvider, just about adding a generic impl<T:Fn()> UrlProvider for T to allow both variants.

/// Provides database url to create a new connection
fn provide_url(&self) -> String;
}

#[cfg(test)]
mod tests {
use std::sync::mpsc;
Expand Down Expand Up @@ -242,4 +266,26 @@ mod tests {
let query = select("foo".into_sql::<Text>());
assert_eq!("foo", query.get_result::<String>(&conn).unwrap());
}

#[derive(Debug)]
struct TestUrlProvider {}
impl UrlProvider for TestUrlProvider {
fn provide_url(&self) -> String {
database_url()
}
}

#[test]
fn provide_dynamic_url() {
let url_provider: Arc<dyn UrlProvider> = Arc::new(TestUrlProvider {});
let manager =
ConnectionManager::<TestConnection>::new_with_url_provider(url_provider.clone());
let pool = Pool::builder()
.max_size(1)
.test_on_check_out(true)
.build(manager)
.unwrap();

pool.get().unwrap();
}
}