opendal/services/onedrive/
writer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use bytes::Buf;
21use bytes::Bytes;
22use http::StatusCode;
23
24use super::core::OneDriveCore;
25use super::error::parse_error;
26use super::graph_model::OneDriveItem;
27use super::graph_model::OneDriveUploadSessionCreationResponseBody;
28use crate::raw::*;
29use crate::*;
30
31pub struct OneDriveWriter {
32    core: Arc<OneDriveCore>,
33    op: OpWrite,
34    path: String,
35}
36
37impl OneDriveWriter {
38    const MAX_SIMPLE_SIZE: usize = 4 * 1024 * 1024; // 4MB
39
40    // OneDrive demands the chunk size to be a multiple to to 320 KiB.
41    // Choose a value smaller than `MAX_SIMPLE_SIZE`
42    const CHUNK_SIZE_FACTOR: usize = 327_680 * 12; // floor(MAX_SIMPLE_SIZE / 320KB)
43
44    pub fn new(core: Arc<OneDriveCore>, op: OpWrite, path: String) -> Self {
45        OneDriveWriter { core, op, path }
46    }
47}
48
49// OneDrive requires multipart writes to specify the total size of the file.
50// While OpenDAL supports multiple writes, due to OneDrive's limitation, we opt to
51// use `OneShotWrite` instead of `MultipartWrite`.
52impl oio::OneShotWrite for OneDriveWriter {
53    async fn write_once(&self, bs: Buffer) -> Result<Metadata> {
54        let size = bs.len();
55
56        let meta = if size <= Self::MAX_SIMPLE_SIZE {
57            self.write_simple(bs).await?
58        } else {
59            self.write_chunked(bs).await?
60        };
61
62        Ok(meta)
63    }
64}
65
66impl OneDriveWriter {
67    async fn write_simple(&self, bs: Buffer) -> Result<Metadata> {
68        let response = self
69            .core
70            .onedrive_upload_simple(&self.path, &self.op, bs)
71            .await?;
72
73        match response.status() {
74            StatusCode::CREATED | StatusCode::OK => {
75                let item: OneDriveItem = serde_json::from_reader(response.into_body().reader())
76                    .map_err(new_json_deserialize_error)?;
77
78                let mut meta = Metadata::new(EntryMode::FILE)
79                    .with_etag(item.e_tag)
80                    .with_content_length(item.size.max(0) as u64);
81
82                let last_modified = item.last_modified_date_time;
83                let date_utc_last_modified = parse_datetime_from_rfc3339(&last_modified)?;
84                meta.set_last_modified(date_utc_last_modified);
85
86                Ok(meta)
87            }
88            _ => Err(parse_error(response)),
89        }
90    }
91
92    pub(crate) async fn write_chunked(&self, bs: Buffer) -> Result<Metadata> {
93        // Upload large files via sessions: https://learn.microsoft.com/en-us/onedrive/developer/rest-api/api/driveitem_createuploadsession?view=odsp-graph-online#upload-bytes-to-the-upload-session
94        // 1. Create an upload session
95        // 2. Upload the bytes of each chunk
96        // 3. Commit the session
97
98        let session_response = self.create_upload_session().await?;
99
100        let mut offset = 0;
101        let total_bytes = bs.to_bytes();
102        let total_len = total_bytes.len();
103        let chunks = total_bytes.chunks(OneDriveWriter::CHUNK_SIZE_FACTOR);
104
105        for chunk in chunks {
106            let mut end = offset + OneDriveWriter::CHUNK_SIZE_FACTOR;
107            if end > total_bytes.len() {
108                end = total_bytes.len();
109            }
110            let chunk_end = end - 1;
111
112            let response = self
113                .core
114                .onedrive_chunked_upload(
115                    &session_response.upload_url,
116                    &self.op,
117                    offset,
118                    chunk_end,
119                    total_len,
120                    Buffer::from(Bytes::copy_from_slice(chunk)),
121                )
122                .await?;
123
124            match response.status() {
125                // Typical response code: 202 Accepted
126                // Reference: https://learn.microsoft.com/en-us/onedrive/developer/rest-api/api/driveitem_put_content?view=odsp-graph-online#response
127                StatusCode::ACCEPTED | StatusCode::OK => {} // skip, in the middle of upload
128                StatusCode::CREATED => {
129                    // last trunk
130                    let item: OneDriveItem = serde_json::from_reader(response.into_body().reader())
131                        .map_err(new_json_deserialize_error)?;
132
133                    let mut meta = Metadata::new(EntryMode::FILE)
134                        .with_etag(item.e_tag)
135                        .with_content_length(item.size.max(0) as u64);
136
137                    let last_modified = item.last_modified_date_time;
138                    let date_utc_last_modified = parse_datetime_from_rfc3339(&last_modified)?;
139                    meta.set_last_modified(date_utc_last_modified);
140                    return Ok(meta);
141                }
142                _ => return Err(parse_error(response)),
143            }
144
145            offset += OneDriveWriter::CHUNK_SIZE_FACTOR;
146        }
147
148        debug_assert!(false, "should have returned");
149
150        Ok(Metadata::default()) // should not happen, but start with handling this gracefully - do nothing, but return the default metadata
151    }
152
153    async fn create_upload_session(&self) -> Result<OneDriveUploadSessionCreationResponseBody> {
154        let response = self
155            .core
156            .onedrive_create_upload_session(&self.path, &self.op)
157            .await?;
158        match response.status() {
159            StatusCode::OK => {
160                let bs = response.into_body();
161                let result: OneDriveUploadSessionCreationResponseBody =
162                    serde_json::from_reader(bs.reader()).map_err(new_json_deserialize_error)?;
163                Ok(result)
164            }
165            _ => Err(parse_error(response)),
166        }
167    }
168}