docbox_core/notifications/
sqs.rs

1use super::{NotificationQueue, NotificationQueueMessage};
2use crate::aws::SqsClient;
3use std::time::Duration;
4use tokio::{spawn, sync::mpsc, time::sleep};
5
6pub struct SqsNotificationQueue {
7    rx: mpsc::Receiver<NotificationQueueMessage>,
8}
9
10impl SqsNotificationQueue {
11    pub fn create(client: SqsClient, queue_url: String) -> SqsNotificationQueue {
12        let (tx, rx) = mpsc::channel(10);
13        let task = SqsNotificationQueueTask {
14            client,
15            queue_url,
16            tx,
17        };
18
19        spawn(process_sqs_queue(task));
20        SqsNotificationQueue { rx }
21    }
22}
23
24impl NotificationQueue for SqsNotificationQueue {
25    async fn next_message(&mut self) -> Option<NotificationQueueMessage> {
26        self.rx.recv().await
27    }
28}
29
30pub struct SqsNotificationQueueTask {
31    /// Underlying client to request messages from
32    client: SqsClient,
33
34    /// URL of the queue containing notifications
35    queue_url: String,
36
37    /// Sender for sending of messages that are ready
38    tx: mpsc::Sender<NotificationQueueMessage>,
39}
40
41pub fn parse_bucket_message(value: &serde_json::Value) -> Option<(String, String)> {
42    let records = value.get("Records")?;
43    let record = records.get(0)?;
44
45    let s3 = record.get("s3")?;
46    let bucket = s3.get("bucket")?;
47    let object = s3.get("object")?;
48
49    let bucket_name = bucket.get("name")?.as_str()?.to_string();
50    let object_key = object.get("key")?.as_str()?.to_string();
51
52    Some((bucket_name, object_key))
53}
54
55async fn process_sqs_queue(task: SqsNotificationQueueTask) {
56    loop {
57        // Receive messages from the SQS queue
58        let receive_messages = match task
59            .client
60            .receive_message()
61            .queue_url(&task.queue_url)
62            .max_number_of_messages(10)
63            .wait_time_seconds(5)
64            .send()
65            .await
66        {
67            Ok(value) => value,
68            Err(error) => {
69                tracing::error!(?error, "error getting messages from sqs");
70                sleep(Duration::from_secs(10)).await;
71                continue;
72            }
73        };
74
75        let messages = match receive_messages.messages {
76            Some(value) => value,
77            None => {
78                tracing::debug!("no messages from sqs");
79                sleep(Duration::from_secs(1)).await;
80                continue;
81            }
82        };
83
84        for message in messages {
85            let (body, receipt_handle) = match (message.body, message.receipt_handle) {
86                (Some(body), Some(receipt_handle)) => (body, receipt_handle),
87                _ => continue,
88            };
89
90            let parsed: serde_json::Value = match serde_json::from_str(&body) {
91                Ok(value) => value,
92                Err(error) => {
93                    if let Err(error) = task
94                        .client
95                        .delete_message()
96                        .queue_url(&task.queue_url)
97                        .receipt_handle(receipt_handle)
98                        .send()
99                        .await
100                    {
101                        tracing::error!(?error, "failed to delete message from sqs");
102                    }
103
104                    tracing::error!(?error, "got malformed message from sqs");
105                    continue;
106                }
107            };
108
109            tracing::debug!(?parsed, "got message from sqs");
110
111            if let Some((bucket_name, object_key)) = parse_bucket_message(&parsed) {
112                tracing::debug!(?bucket_name, ?object_key, "got file upload message");
113
114                _ = task
115                    .tx
116                    .send(NotificationQueueMessage::FileCreated {
117                        bucket_name,
118                        object_key,
119                    })
120                    .await;
121            }
122
123            if let Err(error) = task
124                .client
125                .delete_message()
126                .queue_url(&task.queue_url)
127                .receipt_handle(receipt_handle)
128                .send()
129                .await
130            {
131                tracing::error!(?error, "failed to delete message from sqs");
132            }
133        }
134    }
135}
docbox_core/notifications/sqs.rs

docbox_core/notifications/
sqs.rs
