Extract SqsQueue trait

Author: Turbo87

src/lib.rs

@@ -53,6 +53,7 @@ mod router;
 pub mod schema;
 pub mod sentry;
 pub mod sql;
+pub mod sqs;
 pub mod ssh;
 pub mod storage;
 pub mod tasks;

src/sqs.rs

@@ -0,0 +1,94 @@
+use anyhow::Context;
+use async_trait::async_trait;
+use aws_credential_types::Credentials;
+use aws_sdk_sqs::config::{BehaviorVersion, Region};
+use aws_sdk_sqs::operation::receive_message::ReceiveMessageOutput;
+use mockall::automock;
+
+/// The [SqsQueue] trait defines a basic interface for interacting with an
+/// AWS SQS queue.
+///
+/// A [MockSqsQueue] struct is automatically generated by the [automock]
+/// attribute. This struct can be used in unit tests to mock the behavior of
+/// the [SqsQueue] trait.
+///
+/// The [SqsQueueImpl] struct is the actual implementation of the trait.
+#[automock]
+#[async_trait]
+pub trait SqsQueue {
+    async fn receive_messages(&self, max_messages: i32) -> anyhow::Result<ReceiveMessageOutput>;
+    async fn delete_message(&self, receipt_handle: &str) -> anyhow::Result<()>;
+}
+
+/// The [SqsQueueImpl] struct is the actual implementation of the [SqsQueue]
+/// trait, which interacts with the real AWS API servers.
+#[derive(Debug, Clone)]
+pub struct SqsQueueImpl {
+    client: aws_sdk_sqs::Client,
+    queue_url: String,
+}
+
+impl SqsQueueImpl {
+    pub fn new(queue_url: impl Into<String>, region: Region, credentials: Credentials) -> Self {
+        let config = aws_sdk_sqs::Config::builder()
+            .credentials_provider(credentials)
+            .region(region)
+            .behavior_version(BehaviorVersion::v2023_11_09())
+            .build();
+
+        let client = aws_sdk_sqs::Client::from_conf(config);
+        let queue_url = queue_url.into();
+
+        SqsQueueImpl { client, queue_url }
+    }
+}
+
+#[async_trait]
+impl SqsQueue for SqsQueueImpl {
+    async fn receive_messages(&self, max_messages: i32) -> anyhow::Result<ReceiveMessageOutput> {
+        let response = self
+            .client
+            .receive_message()
+            .max_number_of_messages(max_messages)
+            .queue_url(&self.queue_url)
+            .send()
+            .await
+            .context("Failed to receive SQS queue message")?;
+
+        Ok(response)
+    }
+
+    async fn delete_message(&self, receipt_handle: &str) -> anyhow::Result<()> {
+        self.client
+            .delete_message()
+            .receipt_handle(receipt_handle)
+            .queue_url(&self.queue_url)
+            .send()
+            .await
+            .context("Failed to delete SQS queue message")?;
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_constructor() {
+        let credentials = Credentials::new(
+            "ANOTREAL",
+            "notrealrnrELgWzOk3IfjzDKtFBhDby",
+            None,
+            None,
+            "test",
+        );
+
+        let queue_url = "https://sqs.us-west-1.amazonaws.com/359172468976/cdn-log-event-queue";
+        let region = Region::new("us-west-1");
+
+        // Check that `SqsQueueImpl::new()` does not panic.
+        let _queue = SqsQueueImpl::new(queue_url, region, credentials);
+    }
+}

src/worker/jobs/downloads.rs

@@ -1,11 +1,12 @@
 mod message;
 
 use crate::config::CdnLogStorageConfig;
+use crate::sqs::{SqsQueue, SqsQueueImpl};
 use crate::tasks::spawn_blocking;
 use crate::worker::Environment;
 use anyhow::Context;
 use aws_credential_types::Credentials;
-use aws_sdk_sqs::config::{BehaviorVersion, Region};
+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;
@@ -158,13 +159,7 @@ impl BackgroundJob for ProcessCdnLogQueue {
 
         let credentials = Credentials::from_keys(access_key, secret_key, None);
 
-        let config = aws_sdk_sqs::Config::builder()
-            .credentials_provider(credentials)
-            .region(Region::new(region))
-            .behavior_version(BehaviorVersion::v2023_11_09())
-            .build();
-
-        let client = aws_sdk_sqs::Client::from_conf(config);
+        let queue = SqsQueueImpl::new(queue_url, Region::new(region), credentials);
 
         info!("Receiving messages from the CDN log queue…");
         let mut num_remaining = self.max_messages;
@@ -173,12 +168,7 @@ impl BackgroundJob for ProcessCdnLogQueue {
             num_remaining -= batch_size;
 
             debug!("Receiving next {batch_size} messages from the CDN log queue…");
-            let response = client
-                .receive_message()
-                .queue_url(&queue_url)
-                .max_number_of_messages(batch_size as i32)
-                .send()
-                .await?;
+            let response = queue.receive_messages(batch_size as i32).await?;
 
             let messages = response.messages();
             debug!(
@@ -200,11 +190,8 @@ impl BackgroundJob for ProcessCdnLogQueue {
                 };
 
                 debug!("Deleting message {message_id} from the CDN log queue…");
-                client
-                    .delete_message()
-                    .queue_url(&queue_url)
-                    .receipt_handle(receipt_handle)
-                    .send()
+                queue
+                    .delete_message(receipt_handle)
                     .await
                     .with_context(|| {
                         format!("Failed to delete message {message_id} from the CDN log queue")