1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

#![doc = include_str!("../README.md")]

use tokio::runtime::Runtime;

use async_trait::async_trait;
use shuttle_service::{database, error::CustomError, Error, Factory, ResourceBuilder};

#[cfg(feature = "postgres")]
pub struct Postgres {
    local_uri: Option<String>,
}

#[cfg(feature = "postgres")]
/// Get an `sqlx::PgPool` from any factory
#[async_trait]
impl ResourceBuilder<sqlx::PgPool> for Postgres {
    fn new() -> Self {
        Self { local_uri: None }
    }

    async fn build(
        self,
        factory: &mut dyn Factory,
        runtime: &Runtime,
    ) -> Result<sqlx::PgPool, Error> {
        let connection_string = match factory.get_environment() {
            shuttle_service::Environment::Production => {
                factory
                    .get_db_connection_string(database::Type::Shared(
                        database::SharedEngine::Postgres,
                    ))
                    .await?
            }
            shuttle_service::Environment::Local => {
                if let Some(local_uri) = self.local_uri {
                    local_uri
                } else {
                    factory
                        .get_db_connection_string(database::Type::Shared(
                            database::SharedEngine::Postgres,
                        ))
                        .await?
                }
            }
        };

        // A sqlx Pool cannot cross runtime boundaries, so make sure to create the Pool on the service end
        let pool = runtime
            .spawn(async move {
                sqlx::postgres::PgPoolOptions::new()
                    .min_connections(1)
                    .max_connections(5)
                    .connect(&connection_string)
                    .await
            })
            .await
            .map_err(CustomError::new)?
            .map_err(CustomError::new)?;

        Ok(pool)
    }
}

#[cfg(feature = "postgres")]
impl Postgres {
    /// Use a custom connection string for local runs
    pub fn local_uri(mut self, local_uri: &str) -> Self {
        self.local_uri = Some(local_uri.to_string());

        self
    }
}

#[cfg(feature = "mongodb")]
pub struct MongoDb {
    local_uri: Option<String>,
}

/// Get a `mongodb::Database` from any factory
#[cfg(feature = "mongodb")]
#[async_trait]
impl ResourceBuilder<mongodb::Database> for MongoDb {
    fn new() -> Self {
        Self { local_uri: None }
    }

    async fn build(
        self,
        factory: &mut dyn Factory,
        runtime: &Runtime,
    ) -> Result<mongodb::Database, crate::Error> {
        let connection_string = match factory.get_environment() {
            shuttle_service::Environment::Production => factory
                .get_db_connection_string(database::Type::Shared(database::SharedEngine::MongoDb))
                .await
                .map_err(CustomError::new)?,
            shuttle_service::Environment::Local => {
                if let Some(local_uri) = self.local_uri {
                    local_uri
                } else {
                    factory
                        .get_db_connection_string(database::Type::Shared(
                            database::SharedEngine::MongoDb,
                        ))
                        .await
                        .map_err(CustomError::new)?
                }
            }
        };

        let mut client_options = mongodb::options::ClientOptions::parse(connection_string)
            .await
            .map_err(CustomError::new)?;
        client_options.min_pool_size = Some(1);
        client_options.max_pool_size = Some(5);

        // A mongodb client cannot cross runtime boundaries, so make sure to create the client on the service end
        let client = runtime
            .spawn(async move { mongodb::Client::with_options(client_options) })
            .await
            .map_err(CustomError::new)?
            .map_err(CustomError::new)?;

        // Return a handle to the database defined at the end of the connection string, which is the users provisioned
        // database
        let database = client.default_database();

        match database {
            Some(database) => Ok(database),
            None => Err(crate::Error::Database(
                "mongodb connection string missing default database".into(),
            )),
        }
    }
}

#[cfg(feature = "mongodb")]
impl MongoDb {
    /// Use a custom connection string for local runs
    pub fn local_uri(mut self, local_uri: &str) -> Self {
        self.local_uri = Some(local_uri.to_string());

        self
    }
}