sqs-lambda 0.21.1

SQS Service helper library
Documentation 
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

use std::error::Error;
use std::io::Read;
use std::marker::PhantomData;
use std::time::Duration;

use rusoto_s3::{GetObjectRequest, S3};
use rusoto_sqs::Message as SqsMessage;
use tokio::prelude::*;
use tracing::info;

use async_trait::async_trait;

use crate::event_decoder::PayloadDecoder;
use std::collections::HashMap;

#[async_trait]
pub trait PayloadRetriever<T> {
    type Message;
    async fn retrieve_event(&mut self, msg: &Self::Message) -> Result<Option<T>, Box<dyn Error>>;
}

#[derive(Clone)]
pub struct S3PayloadRetriever<S, SInit, D, E>
where
    S: S3 + Clone + Send + Sync + 'static,
    SInit: (Fn(String) -> S) + Clone + Send + Sync + 'static,
    D: PayloadDecoder<E> + Clone + Send + 'static,
    E: Send + 'static,
{
    s3_init: SInit,
    s3_clients: HashMap<String, S>,
    decoder: D,
    phantom: PhantomData<E>,
}

impl<S, SInit, D, E> S3PayloadRetriever<S, SInit, D, E>
where
    S: S3 + Clone + Send + Sync + 'static,
    SInit: (Fn(String) -> S) + Clone + Send + Sync + 'static,
    D: PayloadDecoder<E> + Clone + Send + 'static,
    E: Send + 'static,
{
    pub fn new(s3: SInit, decoder: D) -> Self {
        Self {
            s3_init: s3,
            s3_clients: HashMap::new(),
            decoder,
            phantom: PhantomData,
        }
    }

    pub fn get_client(&mut self, region: String) -> S {
        match self.s3_clients.get(&region) {
            Some(s3) => s3.clone(),
            None => {
                let client = (self.s3_init)(region.clone());
                self.s3_clients.insert(region.to_string(), client.clone());
                client
            }
        }
    }
}

#[async_trait]
impl<S, SInit, D, E> PayloadRetriever<E> for S3PayloadRetriever<S, SInit, D, E>
where
    S: S3 + Clone + Send + Sync + 'static,
    SInit: (Fn(String) -> S) + Clone + Send + Sync + 'static,
    D: PayloadDecoder<E> + Clone + Send + 'static,
    E: Send + 'static,
{
    type Message = SqsMessage;
    #[tracing::instrument(skip(self, msg))]
    async fn retrieve_event(&mut self, msg: &Self::Message) -> Result<Option<E>, Box<dyn Error>> {
        let body = msg.body.as_ref().unwrap();
        info!("Got body from message: {}", body);
        let event: serde_json::Value = serde_json::from_str(body)?;

        if let Some(Some(event_str)) = event.get("Event").map(serde_json::Value::as_str) {
            if event_str == "s3:TestEvent" {
                return Ok(None);
            }
        }
        let record = &event["Records"][0]["s3"];

        let bucket = record["bucket"]["name"].as_str().expect("bucket name");
        let key = record["object"]["key"].as_str().expect("object key");

        let region = &event["Records"][0]["awsRegion"].as_str().expect("region");
        let s3 = self.get_client(region.to_string());
        let s3_data = s3.get_object(GetObjectRequest {
            bucket: bucket.to_string(),
            key: key.to_string(),
            ..Default::default()
        });

        let s3_data = tokio::time::timeout(Duration::from_secs(5), s3_data).await??;

        let object_size = record["object"]["size"].as_u64().unwrap_or_default();
        let prealloc = if object_size < 1024 {
            1024
        } else {
            object_size as usize
        };

        info!("Retrieved s3 payload with size : {:?}", prealloc);

        let mut body = Vec::with_capacity(prealloc);

        s3_data
            .body
            .expect("Missing S3 body")
            .into_async_read()
            .read_to_end(&mut body)
            .await?;

        info!("Read s3 payload body");
        self.decoder.decode(body).map(Option::from)
    }
}