mz_storage_operators/s3_oneshot_sink/
pgcopy.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use anyhow::anyhow;
11use aws_types::sdk_config::SdkConfig;
12use mz_aws_util::s3_uploader::{
13    CompletedUpload, S3MultiPartUploadError, S3MultiPartUploader, S3MultiPartUploaderConfig,
14};
15use mz_ore::assert_none;
16use mz_ore::cast::CastFrom;
17use mz_pgcopy::{CopyFormatParams, encode_copy_format, encode_copy_format_header};
18use mz_pgrepr::TextEncodeSettings;
19use mz_repr::{GlobalId, RelationDesc, Row};
20use mz_storage_types::sinks::s3_oneshot_sink::S3KeyManager;
21use mz_storage_types::sinks::{S3SinkFormat, S3UploadInfo};
22use tracing::info;
23
24use super::{CopyToParameters, CopyToS3Uploader};
25
26/// Required state to upload batches to S3
27pub(super) struct PgCopyUploader {
28    /// The output description.
29    desc: RelationDesc,
30    /// Params to format the data.
31    format: CopyFormatParams<'static>,
32    /// The index of the current file within the batch.
33    file_index: usize,
34    /// Provides the appropriate bucket and object keys to use for uploads
35    key_manager: S3KeyManager,
36    /// Identifies the batch that files uploaded by this uploader belong to
37    batch: u64,
38    /// The desired file size. A new file upload will be started
39    /// when the size exceeds this amount.
40    max_file_size: u64,
41    /// The aws sdk config.
42    /// This is an option so that we can get an owned value later to move to a
43    /// spawned tokio task.
44    sdk_config: Option<SdkConfig>,
45    /// Multi-part uploader for the current file.
46    /// Keeping the uploader in an `Option` to later take owned value.
47    current_file_uploader: Option<S3MultiPartUploader>,
48    /// Upload parameters.
49    params: CopyToParameters,
50}
51
52impl CopyToS3Uploader for PgCopyUploader {
53    fn new(
54        sdk_config: SdkConfig,
55        connection_details: S3UploadInfo,
56        sink_id: &GlobalId,
57        batch: u64,
58        params: CopyToParameters,
59    ) -> Result<PgCopyUploader, anyhow::Error> {
60        match connection_details.format {
61            S3SinkFormat::PgCopy(format_params) => Ok(PgCopyUploader {
62                desc: connection_details.desc,
63                sdk_config: Some(sdk_config),
64                format: format_params,
65                key_manager: S3KeyManager::new(sink_id, &connection_details.uri),
66                batch,
67                max_file_size: connection_details.max_file_size,
68                file_index: 0,
69                current_file_uploader: None,
70                params,
71            }),
72            S3SinkFormat::Parquet => anyhow::bail!("Expected PgCopy format"),
73        }
74    }
75
76    /// Finishes any remaining in-progress upload.
77    async fn finish(&mut self) -> Result<(), anyhow::Error> {
78        if let Some(uploader) = self.current_file_uploader.take() {
79            // Moving the aws s3 calls onto tokio tasks instead of using timely runtime.
80            let handle =
81                mz_ore::task::spawn(|| "s3_uploader::finish", async { uploader.finish().await });
82            let CompletedUpload {
83                part_count,
84                total_bytes_uploaded,
85                bucket,
86                key,
87            } = handle.await?;
88            info!(
89                "finished upload: bucket {}, key {}, bytes_uploaded {}, parts_uploaded {}",
90                bucket, key, total_bytes_uploaded, part_count
91            );
92        }
93        Ok(())
94    }
95
96    /// Appends the row to the in-progress upload where it is buffered till it reaches the configured
97    /// `part_size_limit` after which the `S3MultiPartUploader` will upload that part. In case it will
98    /// exceed the max file size of the ongoing upload, then a new `S3MultiPartUploader` for a new file will
99    /// be created and the row data will be appended there.
100    async fn append_row(&mut self, row: &Row) -> Result<(), anyhow::Error> {
101        let mut buf: Vec<u8> = vec![];
102        // encode the row and write to temp buffer.
103        encode_copy_format(
104            &self.format,
105            row,
106            self.desc.typ(),
107            &mut buf,
108            TextEncodeSettings::STABLE,
109        )
110        .map_err(|_| anyhow!("error encoding row"))?;
111
112        if self.current_file_uploader.is_none() {
113            self.start_new_file_upload().await?;
114        }
115        let mut uploader = self.current_file_uploader.as_mut().expect("known exists");
116
117        match uploader.buffer_chunk(&buf) {
118            Ok(_) => Ok(()),
119            Err(S3MultiPartUploadError::UploadExceedsMaxFileLimit(_)) => {
120                // Start a multi part upload of next file.
121                self.start_new_file_upload().await?;
122                uploader = self.current_file_uploader.as_mut().expect("known exists");
123                uploader.buffer_chunk(&buf)?;
124                Ok(())
125            }
126            Err(e) => Err(e.into()),
127        }
128    }
129
130    async fn force_new_file(&mut self) -> Result<(), anyhow::Error> {
131        self.start_new_file_upload().await
132    }
133}
134
135impl PgCopyUploader {
136    /// Creates the uploader for the next file and starts the multi part upload.
137    async fn start_new_file_upload(&mut self) -> Result<(), anyhow::Error> {
138        self.finish().await?;
139        assert_none!(self.current_file_uploader);
140
141        self.file_index += 1;
142        let object_key =
143            self.key_manager
144                .data_key(self.batch, self.file_index, self.format.file_extension());
145        let bucket = self.key_manager.bucket.clone();
146        info!("starting upload: bucket {}, key {}", &bucket, &object_key);
147        let sdk_config = self
148            .sdk_config
149            .take()
150            .expect("sdk_config should always be present");
151        let max_file_size = self.max_file_size;
152        // Moving the aws s3 calls onto tokio tasks instead of using timely runtime.
153        let part_size_limit = u64::cast_from(self.params.s3_multipart_part_size_bytes);
154        let handle = mz_ore::task::spawn(|| "s3_uploader::try_new", async move {
155            let uploader = S3MultiPartUploader::try_new(
156                &sdk_config,
157                bucket,
158                object_key,
159                S3MultiPartUploaderConfig {
160                    part_size_limit,
161                    file_size_limit: max_file_size,
162                },
163            )
164            .await;
165            (uploader, sdk_config)
166        });
167        let (uploader, sdk_config) = handle.await;
168        self.sdk_config = Some(sdk_config);
169        let mut uploader = uploader?;
170        if self.format.requires_header() {
171            let mut buf: Vec<u8> = vec![];
172            encode_copy_format_header(&self.format, &self.desc, &mut buf)
173                .map_err(|_| anyhow!("error encoding header"))?;
174            uploader.buffer_chunk(&buf)?;
175        }
176        self.current_file_uploader = Some(uploader);
177        Ok(())
178    }
179}
180
181/// On CI, these tests are enabled by adding the scratch-aws-access plugin
182/// to the `cargo-test` step in `ci/test/pipeline.template.yml` and setting
183/// `MZ_S3_UPLOADER_TEST_S3_BUCKET` in
184/// `ci/test/cargo-test/mzcompose.py`.
185///
186/// For a Materialize developer, to opt in to these tests locally for
187/// development, follow the AWS access guide:
188///
189/// ```text
190/// https://www.notion.so/materialize/AWS-access-5fbd9513dcdc4e11a7591e8caa5f63fe
191/// ```
192///
193/// then running `source src/aws-util/src/setup_test_env_mz.sh`. You will also have
194/// to run `aws sso login` if you haven't recently.
195#[cfg(test)]
196mod tests {
197    use bytesize::ByteSize;
198    use mz_pgcopy::CopyFormatParams;
199    use mz_repr::{ColumnName, Datum, SqlColumnType, SqlRelationType};
200    use uuid::Uuid;
201
202    use super::*;
203
204    fn s3_bucket_path_for_test() -> Option<(String, String)> {
205        let bucket = match std::env::var("MZ_S3_UPLOADER_TEST_S3_BUCKET") {
206            Ok(bucket) => bucket,
207            Err(_) => {
208                if mz_ore::env::is_var_truthy("CI") {
209                    panic!("CI is supposed to run this test but something has gone wrong!");
210                }
211                return None;
212            }
213        };
214
215        let prefix = Uuid::new_v4().to_string();
216        let path = format!("cargo_test/{}/file", prefix);
217        Some((bucket, path))
218    }
219
220    #[mz_ore::test(tokio::test(flavor = "multi_thread"))]
221    #[cfg_attr(coverage, ignore)] // https://github.com/MaterializeInc/database-issues/issues/5586
222    #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `TLS_method` on OS `linux`
223    #[ignore] // TODO: Reenable against minio so it can run locally
224    async fn test_multiple_files() -> Result<(), anyhow::Error> {
225        let sdk_config = mz_aws_util::defaults().load().await;
226        let (bucket, path) = match s3_bucket_path_for_test() {
227            Some(tuple) => tuple,
228            None => return Ok(()),
229        };
230        let sink_id = GlobalId::User(123);
231        let batch = 456;
232        let typ: SqlRelationType = SqlRelationType::new(vec![SqlColumnType {
233            scalar_type: mz_repr::SqlScalarType::String,
234            nullable: true,
235        }]);
236        let column_names = vec![ColumnName::from("col1")];
237        let desc = RelationDesc::new(typ, column_names.into_iter());
238        let mut uploader = PgCopyUploader::new(
239            sdk_config.clone(),
240            S3UploadInfo {
241                uri: format!("s3://{}/{}", bucket, path),
242                // this is only for testing, users will not be able to set value smaller than 16MB.
243                max_file_size: ByteSize::b(6).as_u64(),
244                desc,
245                format: S3SinkFormat::PgCopy(CopyFormatParams::Csv(Default::default())),
246            },
247            &sink_id,
248            batch,
249            CopyToParameters {
250                s3_multipart_part_size_bytes: 10 * 1024 * 1024,
251                arrow_builder_buffer_ratio: 100,
252                parquet_row_group_ratio: 100,
253            },
254        )?;
255        let mut row = Row::default();
256        // Even though this will exceed max_file_size, it should be successfully uploaded in a single file.
257        row.packer().push(Datum::from("1234567"));
258        uploader.append_row(&row).await?;
259
260        // Since the max_file_size is 6B, this row will be uploaded to a new file.
261        row.packer().push(Datum::Null);
262        uploader.append_row(&row).await?;
263
264        row.packer().push(Datum::from("5678"));
265        uploader.append_row(&row).await?;
266
267        uploader.finish().await?;
268
269        // Based on the max_file_size, the uploader should have uploaded two
270        // files, part-0001.csv and part-0002.csv
271        let s3_client = mz_aws_util::s3::new_client(&sdk_config);
272        let first_file = s3_client
273            .get_object()
274            .bucket(bucket.clone())
275            .key(format!(
276                "{}/mz-{}-batch-{:04}-0001.csv",
277                path, sink_id, batch
278            ))
279            .send()
280            .await
281            .unwrap();
282
283        let body = first_file.body.collect().await.unwrap().into_bytes();
284        let expected_body: &[u8] = b"1234567\n";
285        assert_eq!(body, *expected_body);
286
287        let second_file = s3_client
288            .get_object()
289            .bucket(bucket)
290            .key(format!(
291                "{}/mz-{}-batch-{:04}-0002.csv",
292                path, sink_id, batch
293            ))
294            .send()
295            .await
296            .unwrap();
297
298        let body = second_file.body.collect().await.unwrap().into_bytes();
299        let expected_body: &[u8] = b"\n5678\n";
300        assert_eq!(body, *expected_body);
301
302        Ok(())
303    }
304}
mz_storage_operators/s3_oneshot_sink/pgcopy.rs

mz_storage_operators/s3_oneshot_sink/
pgcopy.rs
