Skip to main content

rustack_s3_core/state/
multipart.rs

1//! Multipart upload state management.
2//!
3//! Tracks in-progress multipart uploads and their constituent parts.
4//! Each [`MultipartUpload`] captures the metadata provided at initiation
5//! time and accumulates [`UploadPart`] entries as they are uploaded.
6
7use std::collections::BTreeMap;
8
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11
12use super::object::{ChecksumData, ObjectMetadata, Owner};
13
14/// An in-progress multipart upload.
15///
16/// Created by `CreateMultipartUpload` and completed or aborted later.
17/// Metadata is captured at creation time and applied to the final object
18/// upon completion.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20#[serde(rename_all = "camelCase")]
21pub struct MultipartUpload {
22    /// Unique identifier for this upload.
23    pub upload_id: String,
24    /// The object key that this upload will create.
25    pub key: String,
26    /// When the upload was initiated.
27    pub initiated: DateTime<Utc>,
28    /// The owner who initiated the upload.
29    pub owner: Owner,
30    /// Object metadata captured at `CreateMultipartUpload` time.
31    pub metadata: ObjectMetadata,
32    /// The checksum algorithm requested for this upload (e.g. `CRC32`, `SHA256`).
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub checksum_algorithm: Option<String>,
35    /// Whether the final checksum should be `FULL_OBJECT` or `COMPOSITE`.
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub checksum_type: Option<String>,
38    /// Parts uploaded so far, keyed by part number (1-based).
39    pub parts: BTreeMap<u32, UploadPart>,
40    /// Server-side encryption algorithm for the final object.
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub sse_algorithm: Option<String>,
43    /// KMS key ID for SSE-KMS encryption.
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub sse_kms_key_id: Option<String>,
46    /// The storage class for the final object.
47    pub storage_class: String,
48}
49
50impl MultipartUpload {
51    /// Create a new multipart upload.
52    #[must_use]
53    pub fn new(upload_id: String, key: String, owner: Owner, metadata: ObjectMetadata) -> Self {
54        Self {
55            upload_id,
56            key,
57            initiated: Utc::now(),
58            owner,
59            metadata,
60            checksum_algorithm: None,
61            checksum_type: None,
62            parts: BTreeMap::new(),
63            sse_algorithm: None,
64            sse_kms_key_id: None,
65            storage_class: "STANDARD".to_owned(),
66        }
67    }
68
69    /// Insert or replace a part in this upload.
70    pub fn put_part(&mut self, part: UploadPart) {
71        self.parts.insert(part.part_number, part);
72    }
73
74    /// Get a part by its number.
75    #[must_use]
76    pub fn get_part(&self, part_number: u32) -> Option<&UploadPart> {
77        self.parts.get(&part_number)
78    }
79
80    /// Return the total number of parts uploaded so far.
81    #[must_use]
82    pub fn parts_count(&self) -> usize {
83        self.parts.len()
84    }
85
86    /// Compute the total size of all uploaded parts.
87    #[must_use]
88    pub fn total_size(&self) -> u64 {
89        self.parts.values().map(|p| p.size).sum()
90    }
91}
92
93/// A single part within a multipart upload.
94#[derive(Debug, Clone, Serialize, Deserialize)]
95#[serde(rename_all = "camelCase")]
96pub struct UploadPart {
97    /// The part number (1-based, up to 10 000).
98    pub part_number: u32,
99    /// The entity tag for this part (quoted hex MD5).
100    pub etag: String,
101    /// Size of this part in bytes.
102    pub size: u64,
103    /// When this part was last modified / uploaded.
104    pub last_modified: DateTime<Utc>,
105    /// Optional checksum data for this part.
106    #[serde(skip_serializing_if = "Option::is_none")]
107    pub checksum: Option<ChecksumData>,
108}
109
110// ---------------------------------------------------------------------------
111// Tests
112// ---------------------------------------------------------------------------
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_should_create_multipart_upload() {
120        let upload = MultipartUpload::new(
121            "upload-123".to_owned(),
122            "my-key".to_owned(),
123            Owner::default(),
124            ObjectMetadata::default(),
125        );
126
127        assert_eq!(upload.upload_id, "upload-123");
128        assert_eq!(upload.key, "my-key");
129        assert_eq!(upload.storage_class, "STANDARD");
130        assert_eq!(upload.parts_count(), 0);
131        assert_eq!(upload.total_size(), 0);
132    }
133
134    #[test]
135    fn test_should_put_and_get_parts() {
136        let mut upload = MultipartUpload::new(
137            "upload-456".to_owned(),
138            "data.bin".to_owned(),
139            Owner::default(),
140            ObjectMetadata::default(),
141        );
142
143        let part1 = UploadPart {
144            part_number: 1,
145            etag: "\"abc123\"".to_owned(),
146            size: 5 * 1024 * 1024,
147            last_modified: Utc::now(),
148            checksum: None,
149        };
150        let part2 = UploadPart {
151            part_number: 2,
152            etag: "\"def456\"".to_owned(),
153            size: 3 * 1024 * 1024,
154            last_modified: Utc::now(),
155            checksum: None,
156        };
157
158        upload.put_part(part1);
159        upload.put_part(part2);
160
161        assert_eq!(upload.parts_count(), 2);
162        assert_eq!(upload.total_size(), 8 * 1024 * 1024);
163
164        let p1 = upload.get_part(1);
165        assert!(p1.is_some());
166        assert_eq!(p1.map(|p| &p.etag), Some(&"\"abc123\"".to_owned()));
167
168        assert!(upload.get_part(3).is_none());
169    }
170
171    #[test]
172    fn test_should_replace_existing_part() {
173        let mut upload = MultipartUpload::new(
174            "upload-789".to_owned(),
175            "replace.bin".to_owned(),
176            Owner::default(),
177            ObjectMetadata::default(),
178        );
179
180        let part_v1 = UploadPart {
181            part_number: 1,
182            etag: "\"old\"".to_owned(),
183            size: 100,
184            last_modified: Utc::now(),
185            checksum: None,
186        };
187        upload.put_part(part_v1);
188
189        let part_v2 = UploadPart {
190            part_number: 1,
191            etag: "\"new\"".to_owned(),
192            size: 200,
193            last_modified: Utc::now(),
194            checksum: None,
195        };
196        upload.put_part(part_v2);
197
198        assert_eq!(upload.parts_count(), 1);
199        assert_eq!(upload.total_size(), 200);
200        assert_eq!(
201            upload.get_part(1).map(|p| &p.etag),
202            Some(&"\"new\"".to_owned()),
203        );
204    }
205
206    #[test]
207    fn test_should_store_checksum_on_part() {
208        let part = UploadPart {
209            part_number: 1,
210            etag: "\"abc\"".to_owned(),
211            size: 1024,
212            last_modified: Utc::now(),
213            checksum: Some(ChecksumData {
214                algorithm: "CRC32".to_owned(),
215                value: "AAAAAA==".to_owned(),
216                checksum_type: "FULL_OBJECT".to_owned(),
217            }),
218        };
219        let cs = part.checksum.as_ref();
220        assert!(cs.is_some());
221        assert_eq!(cs.map(|c| c.algorithm.as_str()), Some("CRC32"));
222    }
223}