tracing_s3/
s3_helpers.rs

1use aws_sdk_s3::Client;
2use aws_sdk_s3::primitives::ByteStream;
3use aws_sdk_s3::types::ChecksumAlgorithm;
4
5/// Helper utilities for S3 operations.
6/// Provides methods for file size retrieval and appending content to S3 objects.
7pub struct S3Helpers {}
8
9impl S3Helpers {
10    /// Retrieves the size of a file in S3.
11    /// 
12    /// # Arguments
13    /// * `client` - The AWS S3 client
14    /// * `bucket` - The S3 bucket name
15    /// * `key` - The S3 object key
16    /// 
17    /// # Returns
18    /// * `Ok(i64)` - The file size in bytes, or 0 if the file doesn't exist
19    /// * `Err(anyhow::Error)` - If the head object operation fails
20    pub async fn get_file_size(client: &Client, bucket: &str, key: &str) -> anyhow::Result<i64> {
21        let resp = client.head_object().bucket(bucket).key(key).send().await?;
22        Ok(resp.content_length.unwrap_or(0))
23    }
24    /// Appends content to an existing S3 object or creates a new one if it doesn't exist.
25    /// Uses S3's write_offset_bytes feature for efficient appending.
26    /// 
27    /// # Arguments
28    /// * `client` - The AWS S3 client
29    /// * `bucket` - The S3 bucket name
30    /// * `key` - The S3 object key
31    /// * `content_to_append` - The content to append to the file
32    /// 
33    /// # Returns
34    /// * `Ok(u64)` - The total file size after appending
35    /// * `Err(anyhow::Error)` - If the append operation fails
36    pub async fn append_to_file(
37        client: &Client,
38        bucket: &str,
39        key: &str,
40        content_to_append: &str,
41    ) -> anyhow::Result<u64> {
42        let offset = Self::get_file_size(client, bucket, key).await.unwrap_or(0);
43        let total_len = offset as u64 + content_to_append.len() as u64;
44        let content_to_append = content_to_append.as_bytes().to_vec();
45        client
46            .put_object()
47            .set_write_offset_bytes(Some(offset))
48            .checksum_algorithm(ChecksumAlgorithm::Crc64Nvme)
49            .bucket(bucket)
50            .key(key)
51            .body(ByteStream::from(content_to_append)) // Empty byte array
52            .send()
53            .await?;
54        Ok(total_len)
55    }
56}
57
58#[cfg(test)]
59mod tests {
60
61    use crate::config::tracing_s3_config::TracingS3Config;
62    use crate::config::types::{
63        Bucket, BufferSizeLimitKb, CronIntervalInMs, Endpoint, ObjectSizeLimitMb, Postfix, Prefix,
64    };
65    use crate::s3_helpers::S3Helpers;
66    use chrono::Utc;
67
68    #[tokio::test]
69    pub async fn append_to_file_test() {
70        let config = TracingS3Config::new(
71            None,
72            None,
73            None,
74            Bucket(None),
75            Prefix("prefix"),
76            Postfix("log"),
77            Endpoint(None),
78            ObjectSizeLimitMb::new(1).unwrap(),
79            CronIntervalInMs::new(1_000).unwrap(),
80            BufferSizeLimitKb::new(1).unwrap(),
81        )
82        .await
83        .unwrap();
84        for _ in 0..5 {
85            let _result = S3Helpers::append_to_file(
86                &config.aws_client,
87                &config.bucket,
88                "check-file-exists.log",
89                &format!("hello world {}\n", Utc::now()),
90            )
91            .await;
92        }
93    }
94}