worker/jobs/downloads: Extract CdnLogQueueConfig enum

Author: Turbo87

src/config.rs

@@ -1,12 +1,14 @@
 mod balance_capacity;
 mod base;
+mod cdn_log_queue;
 mod cdn_log_storage;
 mod database_pools;
 mod sentry;
 mod server;
 
 pub use self::balance_capacity::BalanceCapacityConfig;
 pub use self::base::Base;
+pub use self::cdn_log_queue::CdnLogQueueConfig;
 pub use self::cdn_log_storage::CdnLogStorageConfig;
 pub use self::database_pools::{DatabasePools, DbPoolConfig};
 pub use self::sentry::SentryConfig;

src/config/cdn_log_queue.rs

@@ -0,0 +1,33 @@
+use crates_io_env_vars::{required_var, var};
+use secrecy::SecretString;
+
+#[derive(Debug, Clone)]
+pub enum CdnLogQueueConfig {
+    SQS {
+        access_key: String,
+        secret_key: SecretString,
+        queue_url: String,
+        region: String,
+    },
+    Mock,
+}
+
+impl CdnLogQueueConfig {
+    pub fn from_env() -> anyhow::Result<Self> {
+        if let Some(queue_url) = var("CDN_LOG_QUEUE_URL")? {
+            let access_key = required_var("CDN_LOG_QUEUE_ACCESS_KEY")?;
+            let secret_key = required_var("CDN_LOG_QUEUE_SECRET_KEY")?.into();
+            let region = required_var("CDN_LOG_QUEUE_REGION")?;
+
+            return Ok(Self::SQS {
+                access_key,
+                secret_key,
+                queue_url,
+                region,
+            });
+        }
+
+        warn!("Falling back to mocked CDN log queue");
+        Ok(Self::Mock)
+    }
+}

src/config/server.rs

@@ -9,6 +9,7 @@ use super::base::Base;
 use super::database_pools::DatabasePools;
 use crate::config::balance_capacity::BalanceCapacityConfig;
 use crate::config::cdn_log_storage::CdnLogStorageConfig;
+use crate::config::CdnLogQueueConfig;
 use crate::middleware::cargo_compat::StatusCodeConfig;
 use crate::storage::StorageConfig;
 use crates_io_env_vars::{list, list_parsed, required_var, var, var_parsed};
@@ -36,6 +37,7 @@ pub struct Server {
     pub db: DatabasePools,
     pub storage: StorageConfig,
     pub cdn_log_storage: CdnLogStorageConfig,
+    pub cdn_log_queue: CdnLogQueueConfig,
     pub session_key: cookie::Key,
     pub gh_client_id: ClientId,
     pub gh_client_secret: ClientSecret,
@@ -175,6 +177,7 @@ impl Server {
             db: DatabasePools::full_from_environment(&base)?,
             storage,
             cdn_log_storage: CdnLogStorageConfig::from_env()?,
+            cdn_log_queue: CdnLogQueueConfig::from_env()?,
             base,
             ip,
             port,

src/tests/util/test_app.rs

@@ -3,7 +3,8 @@ use crate::util::chaosproxy::ChaosProxy;
 use crate::util::github::{MockGitHubClient, MOCK_GITHUB_DATA};
 use anyhow::Context;
 use crates_io::config::{
-    self, BalanceCapacityConfig, Base, CdnLogStorageConfig, DatabasePools, DbPoolConfig,
+    self, BalanceCapacityConfig, Base, CdnLogQueueConfig, CdnLogStorageConfig, DatabasePools,
+    DbPoolConfig,
 };
 use crates_io::middleware::cargo_compat::StatusCodeConfig;
 use crates_io::models::token::{CrateScope, EndpointScope};
@@ -424,6 +425,7 @@ fn simple_config() -> config::Server {
         max_blocking_threads: None,
         db,
         storage,
+        cdn_log_queue: CdnLogQueueConfig::Mock,
         cdn_log_storage: CdnLogStorageConfig::memory(),
         session_key: cookie::Key::derive_from("test this has to be over 32 bytes long".as_bytes()),
         gh_client_id: ClientId::new(dotenvy::var("GH_CLIENT_ID").unwrap_or_default()),

src/worker/jobs/downloads.rs

@@ -1,14 +1,13 @@
 mod message;
 
-use crate::config::CdnLogStorageConfig;
-use crate::sqs::{SqsQueue, SqsQueueImpl};
+use crate::config::{CdnLogQueueConfig, CdnLogStorageConfig};
+use crate::sqs::{MockSqsQueue, SqsQueue, SqsQueueImpl};
 use crate::tasks::spawn_blocking;
 use crate::worker::Environment;
 use anyhow::Context;
 use aws_credential_types::Credentials;
 use aws_sdk_sqs::config::Region;
 use crates_io_cdn_logs::{count_downloads, Decompressor};
-use crates_io_env_vars::required_var;
 use crates_io_worker::BackgroundJob;
 use object_store::aws::AmazonS3Builder;
 use object_store::local::LocalFileSystem;
@@ -152,14 +151,7 @@ impl BackgroundJob for ProcessCdnLogQueue {
     async fn run(&self, ctx: Self::Context) -> anyhow::Result<()> {
         const MAX_BATCH_SIZE: usize = 10;
 
-        let access_key = required_var("CDN_LOG_QUEUE_ACCESS_KEY")?;
-        let secret_key = required_var("CDN_LOG_QUEUE_SECRET_KEY")?;
-        let queue_url = required_var("CDN_LOG_QUEUE_URL")?;
-        let region = required_var("CDN_LOG_QUEUE_REGION")?;
-
-        let credentials = Credentials::from_keys(access_key, secret_key, None);
-
-        let queue = SqsQueueImpl::new(queue_url, Region::new(region), credentials);
+        let queue = Self::build_queue(&ctx.config.cdn_log_queue);
 
         info!("Receiving messages from the CDN log queue…");
         let mut num_remaining = self.max_messages;
@@ -259,6 +251,29 @@ impl BackgroundJob for ProcessCdnLogQueue {
     }
 }
 
+impl ProcessCdnLogQueue {
+    fn build_queue(config: &CdnLogQueueConfig) -> Box<dyn SqsQueue + Send + Sync> {
+        match config {
+            CdnLogQueueConfig::Mock => Box::new(MockSqsQueue::new()),
+            CdnLogQueueConfig::SQS {
+                access_key,
+                secret_key,
+                region,
+                queue_url,
+            } => {
+                use secrecy::ExposeSecret;
+
+                let secret_key = secret_key.expose_secret();
+                let credentials = Credentials::from_keys(access_key, secret_key, None);
+
+                let region = Region::new(region.to_owned());
+
+                Box::new(SqsQueueImpl::new(queue_url, region, credentials))
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;