Skip to main content

docbox_database/
pool.rs

1//! # Database Pool
2//!
3//! This is the docbox solution for managing multiple database connections
4//! and connection pools for each tenant and the root database itself.
5//!
6//! Pools are held in a cache with an expiry time to ensure they don't
7//! hog too many database connections.
8//!
9//! Database pools and credentials are stored in a Tiny LFU cache these caches
10//! can be flushed using [DatabasePoolCache::flush]
11//!
12//! ## Environment Variables
13//!
14//! * `DOCBOX_DB_HOST` - Database host
15//! * `DOCBOX_DB_PORT` - Database port
16//! * `DOCBOX_DB_CREDENTIAL_NAME` - Secrets manager name for the root database secret
17//! * `DOCBOX_DB_ROOT_IAM` - Whether to use IAM to authenticate the root database
18//! * `DOCBOX_DB_MAX_CONNECTIONS` - Max connections each tenant pool can contain
19//! * `DOCBOX_DB_MAX_ROOT_CONNECTIONS` - Max connections the root "docbox" pool can contain
20//! * `DOCBOX_DB_ACQUIRE_TIMEOUT` - Timeout before acquiring a connection fails
21//! * `DOCBOX_DB_POOL_TIMEOUT` - Maximum time a connection can live in the cache for
22//! * `DOCBOX_DB_IDLE_TIMEOUT` - Timeout before a idle connection is closed to save resources
23//! * `DOCBOX_DB_CACHE_DURATION` - Duration pools can remain in the cache for untouched before they are closed and removed
24//! * `DOCBOX_DB_CACHE_CAPACITY` - Maximum database pools to hold at once
25//! * `DOCBOX_DB_CREDENTIALS_CACHE_DURATION` - Duration database credentials should be cached for
26//! * `DOCBOX_DB_CREDENTIALS_CACHE_CAPACITY` - Maximum database credentials to cache
27
28use crate::{DbErr, DbPool, ROOT_DATABASE_NAME, ROOT_DATABASE_ROLE_NAME, models::tenant::Tenant};
29use aws_config::SdkConfig;
30use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError};
31use aws_sigv4::{
32    http_request::{SignableBody, SignableRequest, SigningError, SigningSettings, sign},
33    sign::v4::signing_params,
34};
35use docbox_secrets::{SecretManager, SecretManagerError};
36use moka::{future::Cache, policy::EvictionPolicy};
37use serde::{Deserialize, Serialize};
38use sqlx::{
39    PgPool,
40    postgres::{PgConnectOptions, PgPoolOptions},
41};
42use std::time::Duration;
43use std::{num::ParseIntError, str::ParseBoolError};
44use std::{sync::Arc, time::SystemTime};
45use thiserror::Error;
46use tokio::time::sleep;
47
48///  Config for the database pool
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct DatabasePoolCacheConfig {
51    /// Database host
52    pub host: String,
53    /// Database port
54    pub port: u16,
55
56    /// Name of the secrets manager secret to use when connecting to
57    /// the root "docbox" database if using secret based authentication
58    pub root_secret_name: Option<String>,
59
60    /// Whether to use IAM authentication to connect to the
61    /// root database instead of secrets
62    #[serde(default)]
63    pub root_iam: bool,
64
65    /// Max number of active connections per tenant database pool
66    ///
67    /// This is the maximum number of connections that should be allocated
68    /// for performing all queries against each specific tenant.
69    ///
70    /// Ensure a reasonable amount of connections are allocated but make
71    /// sure that the `max_connections` * your number of tenants stays
72    /// within the limits for your database
73    ///
74    /// Default: 10
75    pub max_connections: Option<u32>,
76
77    /// Max number of active connections per "docbox" database pool
78    ///
79    /// This is the maximum number of connections that should be allocated
80    /// for performing queries like:
81    /// - Listing tenants
82    /// - Getting tenant details
83    ///
84    /// These pools are often short lived and complete their queries very fast
85    /// and thus don't need a huge amount of resources allocated to them
86    ///
87    /// Default: 2
88    pub max_connections_root: Option<u32>,
89
90    /// Timeout before a acquiring a database connection is considered
91    /// a failure
92    ///
93    /// Default: 60s
94    pub acquire_timeout: Option<u64>,
95
96    /// If a connection has been idle for this duration the connection
97    /// will be closed and released back to the database for other
98    /// consumers
99    ///
100    /// Default: 10min
101    pub idle_timeout: Option<u64>,
102
103    /// Maximum time pool are allowed to stay within the database
104    /// cache before they are automatically removed
105    ///
106    /// Default: 48h
107    pub pool_timeout: Option<u64>,
108
109    /// Duration in seconds idle database pools are allowed to be cached before
110    /// they are closed
111    ///
112    /// Default: 48h
113    pub cache_duration: Option<u64>,
114
115    /// Maximum database pools to maintain in the cache at once. If the
116    /// cache capacity is exceeded old pools will be closed and removed
117    /// from the cache
118    ///
119    /// This capacity should be aligned with your expected number of
120    /// tenants along with your `max_connections` to ensure your database
121    /// has enough connections to accommodate all tenants.
122    ///
123    /// Default: 50
124    pub cache_capacity: Option<u64>,
125
126    /// Duration in seconds database credentials (host, port, password, ..etc)
127    /// are allowed to be cached before they are refresh from the secrets
128    /// manager
129    ///
130    /// Default: 12h
131    pub credentials_cache_duration: Option<u64>,
132
133    /// Maximum database credentials to maintain in the cache at once. If the
134    /// cache capacity is exceeded old credentials will be removed from the cache
135    ///
136    /// Default: 50
137    pub credentials_cache_capacity: Option<u64>,
138}
139
140impl Default for DatabasePoolCacheConfig {
141    fn default() -> Self {
142        Self {
143            host: Default::default(),
144            port: 5432,
145            root_secret_name: Default::default(),
146            root_iam: false,
147            max_connections: None,
148            max_connections_root: None,
149            acquire_timeout: None,
150            idle_timeout: None,
151            pool_timeout: None,
152            cache_duration: None,
153            cache_capacity: None,
154            credentials_cache_duration: None,
155            credentials_cache_capacity: None,
156        }
157    }
158}
159
160#[derive(Debug, Error)]
161pub enum DatabasePoolCacheConfigError {
162    #[error("missing DOCBOX_DB_HOST environment variable")]
163    MissingDatabaseHost,
164    #[error("missing DOCBOX_DB_PORT environment variable")]
165    MissingDatabasePort,
166    #[error("invalid DOCBOX_DB_PORT environment variable")]
167    InvalidDatabasePort,
168    #[error("missing DOCBOX_DB_CREDENTIAL_NAME environment variable")]
169    MissingDatabaseSecretName,
170    #[error("invalid DOCBOX_DB_POOL_TIMEOUT environment variable")]
171    InvalidPoolTimeout(ParseIntError),
172    #[error("invalid DOCBOX_DB_IDLE_TIMEOUT environment variable")]
173    InvalidIdleTimeout(ParseIntError),
174    #[error("invalid DOCBOX_DB_ACQUIRE_TIMEOUT environment variable")]
175    InvalidAcquireTimeout(ParseIntError),
176    #[error("invalid DOCBOX_DB_CACHE_DURATION environment variable")]
177    InvalidCacheDuration(ParseIntError),
178    #[error("invalid DOCBOX_DB_CACHE_CAPACITY environment variable")]
179    InvalidCacheCapacity(ParseIntError),
180    #[error("invalid DOCBOX_DB_CREDENTIALS_CACHE_DURATION environment variable")]
181    InvalidCredentialsCacheDuration(ParseIntError),
182    #[error("invalid DOCBOX_DB_CREDENTIALS_CACHE_CAPACITY environment variable")]
183    InvalidCredentialsCacheCapacity(ParseIntError),
184    #[error("invalid DOCBOX_DB_ROOT_IAM environment variable")]
185    InvalidRootIam(ParseBoolError),
186}
187
188impl DatabasePoolCacheConfig {
189    pub fn from_env() -> Result<DatabasePoolCacheConfig, DatabasePoolCacheConfigError> {
190        let db_host: String = std::env::var("DOCBOX_DB_HOST")
191            .or(std::env::var("POSTGRES_HOST"))
192            .map_err(|_| DatabasePoolCacheConfigError::MissingDatabaseHost)?;
193        let db_port: u16 = std::env::var("DOCBOX_DB_PORT")
194            .or(std::env::var("POSTGRES_PORT"))
195            .map_err(|_| DatabasePoolCacheConfigError::MissingDatabasePort)?
196            .parse()
197            .map_err(|_| DatabasePoolCacheConfigError::InvalidDatabasePort)?;
198
199        let db_root_secret_name = std::env::var("DOCBOX_DB_CREDENTIAL_NAME").ok();
200        let db_root_iam = std::env::var("DOCBOX_DB_ROOT_IAM")
201            .ok()
202            .map(|value| value.parse::<bool>())
203            .transpose()
204            .map_err(DatabasePoolCacheConfigError::InvalidRootIam)?
205            .unwrap_or_default();
206
207        // Root secret name is required when not using IAM
208        if !db_root_iam && db_root_secret_name.is_none() {
209            return Err(DatabasePoolCacheConfigError::MissingDatabaseSecretName);
210        }
211
212        let max_connections: Option<u32> = std::env::var("DOCBOX_DB_MAX_CONNECTIONS")
213            .ok()
214            .and_then(|value| value.parse().ok());
215        let max_connections_root: Option<u32> = std::env::var("DOCBOX_DB_MAX_ROOT_CONNECTIONS")
216            .ok()
217            .and_then(|value| value.parse().ok());
218
219        let acquire_timeout: Option<u64> = match std::env::var("DOCBOX_DB_ACQUIRE_TIMEOUT") {
220            Ok(value) => Some(
221                value
222                    .parse::<u64>()
223                    .map_err(DatabasePoolCacheConfigError::InvalidAcquireTimeout)?,
224            ),
225            Err(_) => None,
226        };
227
228        let pool_timeout: Option<u64> = match std::env::var("DOCBOX_DB_POOL_TIMEOUT") {
229            Ok(value) => Some(
230                value
231                    .parse::<u64>()
232                    .map_err(DatabasePoolCacheConfigError::InvalidPoolTimeout)?,
233            ),
234            Err(_) => None,
235        };
236
237        let idle_timeout: Option<u64> = match std::env::var("DOCBOX_DB_IDLE_TIMEOUT") {
238            Ok(value) => Some(
239                value
240                    .parse::<u64>()
241                    .map_err(DatabasePoolCacheConfigError::InvalidIdleTimeout)?,
242            ),
243            Err(_) => None,
244        };
245
246        let cache_duration: Option<u64> = match std::env::var("DOCBOX_DB_CACHE_DURATION") {
247            Ok(value) => Some(
248                value
249                    .parse::<u64>()
250                    .map_err(DatabasePoolCacheConfigError::InvalidCacheDuration)?,
251            ),
252            Err(_) => None,
253        };
254
255        let cache_capacity: Option<u64> = match std::env::var("DOCBOX_DB_CACHE_CAPACITY") {
256            Ok(value) => Some(
257                value
258                    .parse::<u64>()
259                    .map_err(DatabasePoolCacheConfigError::InvalidCacheCapacity)?,
260            ),
261            Err(_) => None,
262        };
263
264        let credentials_cache_duration: Option<u64> =
265            match std::env::var("DOCBOX_DB_CREDENTIALS_CACHE_DURATION") {
266                Ok(value) => Some(
267                    value
268                        .parse::<u64>()
269                        .map_err(DatabasePoolCacheConfigError::InvalidCredentialsCacheDuration)?,
270                ),
271                Err(_) => None,
272            };
273
274        let credentials_cache_capacity: Option<u64> =
275            match std::env::var("DOCBOX_DB_CREDENTIALS_CACHE_CAPACITY") {
276                Ok(value) => Some(
277                    value
278                        .parse::<u64>()
279                        .map_err(DatabasePoolCacheConfigError::InvalidCredentialsCacheCapacity)?,
280                ),
281                Err(_) => None,
282            };
283
284        Ok(DatabasePoolCacheConfig {
285            host: db_host,
286            port: db_port,
287            root_iam: db_root_iam,
288            root_secret_name: db_root_secret_name,
289            max_connections,
290            max_connections_root,
291            acquire_timeout,
292            pool_timeout,
293            idle_timeout,
294            cache_duration,
295            cache_capacity,
296            credentials_cache_duration,
297            credentials_cache_capacity,
298        })
299    }
300}
301
302/// Cache for database pools
303pub struct DatabasePoolCache {
304    /// AWS config
305    aws_config: aws_config::SdkConfig,
306
307    /// Database host
308    host: String,
309
310    /// Database port
311    port: u16,
312
313    /// Name of the secrets manager secret that contains
314    /// the credentials for the root "docbox" database
315    ///
316    /// Only present if using secrets based authentication
317    root_secret_name: Option<String>,
318
319    /// Whether to use IAM authentication to connect to the
320    /// root database instead of secrets
321    root_iam: bool,
322
323    /// Cache from the database name to the pool for that database
324    cache: Cache<String, DbPool>,
325
326    /// Cache for the connection info details, stores the last known
327    /// credentials and the instant that they were obtained at
328    connect_info_cache: Cache<String, DbSecrets>,
329
330    /// Secrets manager access to load credentials
331    secrets_manager: SecretManager,
332
333    /// Max connections per tenant database pool
334    max_connections: u32,
335    /// Max connections per root database pool
336    max_connections_root: u32,
337
338    acquire_timeout: Duration,
339    idle_timeout: Duration,
340}
341
342/// Username and password for a specific database
343#[derive(Debug, Clone, Serialize, Deserialize)]
344pub struct DbSecrets {
345    pub username: String,
346    pub password: String,
347}
348
349#[derive(Debug, Error)]
350pub enum DbConnectErr {
351    #[error("database credentials not found in secrets manager")]
352    MissingCredentials,
353
354    #[error(transparent)]
355    SecretsManager(Box<SecretManagerError>),
356
357    #[error(transparent)]
358    Db(#[from] DbErr),
359
360    #[error(transparent)]
361    Shared(#[from] Arc<DbConnectErr>),
362
363    #[error("missing aws credentials provider")]
364    MissingCredentialsProvider,
365
366    #[error("failed to provide aws credentials")]
367    AwsCredentials(#[from] CredentialsError),
368
369    #[error("aws configuration missing region")]
370    MissingRegion,
371
372    #[error("failed to build aws signature")]
373    AwsSigner(#[from] signing_params::BuildError),
374
375    #[error("failed to sign aws request")]
376    AwsRequestSign(#[from] SigningError),
377
378    #[error("failed to parse signed aws url")]
379    AwsSignerInvalidUrl(url::ParseError),
380
381    #[error("failed to connect to tenant missing both IAM and secrets fields")]
382    InvalidTenantConfiguration,
383}
384
385impl DatabasePoolCache {
386    pub fn from_config(
387        aws_config: aws_config::SdkConfig,
388        config: DatabasePoolCacheConfig,
389        secrets_manager: SecretManager,
390    ) -> Self {
391        let mut pool_timeout = Duration::from_secs(config.cache_duration.unwrap_or(60 * 60 * 48));
392        let cache_duration = Duration::from_secs(config.cache_duration.unwrap_or(60 * 60 * 48));
393        let credentials_cache_duration =
394            Duration::from_secs(config.credentials_cache_duration.unwrap_or(60 * 60 * 12));
395
396        // When using IAM ensure the pool timeout is less than the expiration time
397        // of the temporary access tokens
398        if config.root_iam && config.pool_timeout.is_none() {
399            tracing::debug!(
400                "IAM database auth is enabled with no pool timeout, setting short pool timeout within token duration"
401            );
402            pool_timeout = Duration::from_secs(60 * 10);
403        }
404
405        let cache_capacity = config.cache_capacity.unwrap_or(50);
406        let credentials_cache_capacity = config.credentials_cache_capacity.unwrap_or(50);
407
408        let cache = Cache::builder()
409            .time_to_live(pool_timeout)
410            .time_to_idle(cache_duration)
411            .max_capacity(cache_capacity)
412            .eviction_policy(EvictionPolicy::tiny_lfu())
413            .async_eviction_listener(|cache_key: Arc<String>, pool: DbPool, _cause| {
414                Box::pin(async move {
415                    tracing::debug!(?cache_key, "database pool is no longer in use, closing");
416                    pool.close().await
417                })
418            })
419            .build();
420
421        let connect_info_cache = Cache::builder()
422            .time_to_idle(credentials_cache_duration)
423            .max_capacity(credentials_cache_capacity)
424            .eviction_policy(EvictionPolicy::tiny_lfu())
425            .build();
426
427        Self {
428            aws_config,
429            host: config.host,
430            port: config.port,
431            root_secret_name: config.root_secret_name,
432            root_iam: config.root_iam,
433            cache,
434            connect_info_cache,
435            secrets_manager,
436            max_connections: config.max_connections.unwrap_or(10),
437            max_connections_root: config.max_connections_root.unwrap_or(2),
438            idle_timeout: Duration::from_secs(config.idle_timeout.unwrap_or(60 * 10)),
439            acquire_timeout: Duration::from_secs(config.acquire_timeout.unwrap_or(60)),
440        }
441    }
442
443    /// Request a database pool for the root database
444    pub async fn get_root_pool(&self) -> Result<PgPool, DbConnectErr> {
445        match (self.root_secret_name.as_ref(), self.root_iam) {
446            (_, true) => {
447                self.get_pool_iam(ROOT_DATABASE_NAME, ROOT_DATABASE_ROLE_NAME)
448                    .await
449            }
450
451            (Some(db_secret_name), _) => self.get_pool(ROOT_DATABASE_NAME, db_secret_name).await,
452
453            _ => Err(DbConnectErr::InvalidTenantConfiguration),
454        }
455    }
456
457    /// Request a database pool for a specific tenant
458    pub async fn get_tenant_pool(&self, tenant: &Tenant) -> Result<DbPool, DbConnectErr> {
459        match (
460            tenant.db_iam_user_name.as_ref(),
461            tenant.db_secret_name.as_ref(),
462        ) {
463            (Some(db_iam_user_name), _) => {
464                self.get_pool_iam(&tenant.db_name, db_iam_user_name).await
465            }
466            (_, Some(db_secret_name)) => self.get_pool(&tenant.db_name, db_secret_name).await,
467
468            _ => Err(DbConnectErr::InvalidTenantConfiguration),
469        }
470    }
471
472    /// Closes the database pool for the specific tenant if one is
473    /// available and removes the pool from the cache
474    pub async fn close_tenant_pool(&self, tenant: &Tenant) {
475        let cache_key = Self::tenant_cache_key(tenant);
476        if let Some(pool) = self.cache.remove(&cache_key).await {
477            pool.close().await;
478        }
479
480        // Run cache async shutdown jobs
481        self.cache.run_pending_tasks().await;
482    }
483
484    /// Compute the pool cache key for a tenant based on the specific
485    /// authentication methods for that tenant
486    fn tenant_cache_key(tenant: &Tenant) -> String {
487        match (
488            tenant.db_secret_name.as_ref(),
489            tenant.db_iam_user_name.as_ref(),
490        ) {
491            (Some(db_secret_name), _) => {
492                format!("secret-{}-{}", &tenant.db_name, db_secret_name)
493            }
494            (_, Some(db_iam_user_name)) => {
495                format!("user-{}-{}", &tenant.db_name, db_iam_user_name)
496            }
497
498            _ => format!("db-{}", &tenant.db_name),
499        }
500    }
501
502    /// Empties all the caches
503    pub async fn flush(&self) {
504        // Clear cache
505        self.cache.invalidate_all();
506        self.connect_info_cache.invalidate_all();
507        self.cache.run_pending_tasks().await;
508    }
509
510    /// Close all connections in the pool and invalidate the cache
511    pub async fn close_all(&self) {
512        for (_, value) in self.cache.iter() {
513            value.close().await;
514        }
515
516        self.flush().await;
517    }
518
519    /// Obtains a database pool connection to the database with the provided name
520    /// using secrets manager based credentials
521    async fn get_pool(&self, db_name: &str, secret_name: &str) -> Result<DbPool, DbConnectErr> {
522        let cache_key = format!("secret-{db_name}-{secret_name}");
523
524        let pool = self
525            .cache
526            .try_get_with(cache_key, async {
527                tracing::debug!(?db_name, "acquiring database pool");
528
529                let pool = self
530                    .create_pool(db_name, secret_name)
531                    .await
532                    .map_err(Arc::new)?;
533
534                Ok(pool)
535            })
536            .await?;
537
538        Ok(pool)
539    }
540
541    /// Obtains a database pool connection to the database with the provided name
542    /// using IAM based credentials
543    async fn get_pool_iam(
544        &self,
545        db_name: &str,
546        db_role_name: &str,
547    ) -> Result<DbPool, DbConnectErr> {
548        let cache_key = format!("user-{db_name}-{db_role_name}");
549
550        let pool = self
551            .cache
552            .try_get_with(cache_key, async {
553                tracing::debug!(?db_name, "acquiring database pool (iam)");
554
555                let pool = self
556                    .create_pool_iam(db_name, db_role_name)
557                    .await
558                    .map_err(Arc::new)?;
559
560                Ok(pool)
561            })
562            .await?;
563
564        Ok(pool)
565    }
566
567    /// Obtains database connection info
568    async fn get_credentials(&self, secret_name: &str) -> Result<DbSecrets, DbConnectErr> {
569        if let Some(connect_info) = self.connect_info_cache.get(secret_name).await {
570            return Ok(connect_info);
571        }
572
573        // Load new credentials
574        let credentials = self
575            .secrets_manager
576            .parsed_secret::<DbSecrets>(secret_name)
577            .await
578            .map_err(|err| DbConnectErr::SecretsManager(Box::new(err)))?
579            .ok_or(DbConnectErr::MissingCredentials)?;
580
581        // Cache the credential
582        self.connect_info_cache
583            .insert(secret_name.to_string(), credentials.clone())
584            .await;
585
586        Ok(credentials)
587    }
588
589    /// Creates a database pool connection using IAM based authentication
590    async fn create_pool_iam(
591        &self,
592        db_name: &str,
593        db_role_name: &str,
594    ) -> Result<DbPool, DbConnectErr> {
595        tracing::debug!(?db_name, ?db_role_name, "creating db pool connection");
596
597        let options = iam_pool_connect_options(
598            &self.aws_config,
599            &self.host,
600            self.port,
601            db_name,
602            db_role_name,
603        )
604        .await?;
605
606        let max_connections = match db_name {
607            ROOT_DATABASE_NAME => self.max_connections_root,
608            _ => self.max_connections,
609        };
610
611        let pool = PgPoolOptions::new()
612            .max_connections(max_connections)
613            // Slightly larger acquire timeout for times when lots of files are being processed
614            .acquire_timeout(self.acquire_timeout)
615            // Close any connections that have been idle for more than 30min
616            .idle_timeout(self.idle_timeout)
617            .connect_with(options)
618            .await
619            .map_err(DbConnectErr::Db)?;
620
621        tokio::spawn(iam_pool_maintenance_task(
622            pool.clone(),
623            self.aws_config.clone(),
624            self.host.clone(),
625            self.port,
626            db_name.to_string(),
627            db_role_name.to_string(),
628        ));
629
630        Ok(pool)
631    }
632
633    /// Creates a database pool connection
634    async fn create_pool(&self, db_name: &str, secret_name: &str) -> Result<DbPool, DbConnectErr> {
635        tracing::debug!(?db_name, ?secret_name, "creating db pool connection");
636
637        let credentials = self.get_credentials(secret_name).await?;
638        let options = PgConnectOptions::new()
639            .host(&self.host)
640            .port(self.port)
641            .username(&credentials.username)
642            .password(&credentials.password)
643            .database(db_name);
644
645        let max_connections = match db_name {
646            ROOT_DATABASE_NAME => self.max_connections_root,
647            _ => self.max_connections,
648        };
649
650        match PgPoolOptions::new()
651            .max_connections(max_connections)
652            // Slightly larger acquire timeout for times when lots of files are being processed
653            .acquire_timeout(self.acquire_timeout)
654            // Close any connections that have been idle for more than 30min
655            .idle_timeout(self.idle_timeout)
656            .connect_with(options)
657            .await
658        {
659            // Success case
660            Ok(value) => Ok(value),
661            Err(err) => {
662                // Drop the connect info cache in case the credentials were wrong
663                self.connect_info_cache.remove(secret_name).await;
664                Err(DbConnectErr::Db(err))
665            }
666        }
667    }
668}
669
670async fn iam_pool_connect_options(
671    aws_config: &SdkConfig,
672    host: &str,
673    port: u16,
674    db_name: &str,
675    db_role_name: &str,
676) -> Result<PgConnectOptions, DbConnectErr> {
677    let token = create_rds_signed_token(aws_config, host, port, db_role_name).await?;
678
679    let options = PgConnectOptions::new()
680        .host(host)
681        .port(port)
682        .username(db_role_name)
683        .password(&token)
684        .database(db_name);
685
686    Ok(options)
687}
688
689async fn create_rds_signed_token(
690    aws_config: &SdkConfig,
691    host: &str,
692    port: u16,
693    user: &str,
694) -> Result<String, DbConnectErr> {
695    let credentials_provider = aws_config
696        .credentials_provider()
697        .ok_or(DbConnectErr::MissingCredentialsProvider)?;
698    let credentials = credentials_provider.provide_credentials().await?;
699    let identity = credentials.into();
700    let region = aws_config.region().ok_or(DbConnectErr::MissingRegion)?;
701
702    let mut signing_settings = SigningSettings::default();
703    signing_settings.expires_in = Some(Duration::from_secs(60 * 15));
704    signing_settings.signature_location = aws_sigv4::http_request::SignatureLocation::QueryParams;
705
706    let signing_params = aws_sigv4::sign::v4::SigningParams::builder()
707        .identity(&identity)
708        .region(region.as_ref())
709        .name("rds-db")
710        .time(SystemTime::now())
711        .settings(signing_settings)
712        .build()?;
713
714    let url = format!("https://{host}:{port}/?Action=connect&DBUser={user}");
715
716    let signable_request =
717        SignableRequest::new("GET", &url, std::iter::empty(), SignableBody::Bytes(&[]))?;
718
719    let (signing_instructions, _signature) =
720        sign(signable_request, &signing_params.into())?.into_parts();
721
722    let mut url = url::Url::parse(&url).map_err(DbConnectErr::AwsSignerInvalidUrl)?;
723    for (name, value) in signing_instructions.params() {
724        url.query_pairs_mut().append_pair(name, value);
725    }
726
727    let response = url.to_string().split_off("https://".len());
728    Ok(response)
729}
730
731/// Background task spawned for IAM pools running every 10minutes to ensure that the pool
732/// has an up-to-date temporary authentication token
733async fn iam_pool_maintenance_task(
734    db: DbPool,
735    aws_config: SdkConfig,
736    host: String,
737    port: u16,
738    db_name: String,
739    db_role_name: String,
740) {
741    let interval = Duration::from_secs(60 * 10);
742
743    loop {
744        if db.is_closed() {
745            return;
746        }
747
748        match iam_pool_connect_options(&aws_config, &host, port, &db_name, &db_role_name).await {
749            Ok(options) => {
750                db.set_connect_options(options);
751            }
752            Err(error) => {
753                tracing::error!(?error, "failed to refresh IAM pool connect options");
754            }
755        }
756
757        sleep(interval).await;
758    }
759}