drive-v3 0.5.1

A library for interacting the Google Drive API
Documentation
use std::fs;
use std::io::Read;
use reqwest::blocking::Client;

use crate::{Error, ErrorKind, objects};

// https://developers.google.com/drive/api/guides/manage-uploads#http---multiple-requests
const UPLOAD_CHUNK_SIZE: usize =  256 * 1024 * 38; // approx 10Mb

/// Uploads a file in chunks
#[doc(hidden)]
pub struct FileUploader {
    upload_uri: String,
    current_byte: usize,
    total_bytes: usize,
    chunk_size: usize,
    callback: Option<fn(usize, usize)>,
}

impl FileUploader {
    pub fn from_uri<T: AsRef<str>> ( resumable_session_uri: T ) -> Self {
        Self {
            upload_uri: resumable_session_uri.as_ref().to_string(),
            current_byte: 0,
            total_bytes: 0,
            chunk_size: UPLOAD_CHUNK_SIZE,
            callback: None,
        }
    }

    /// Sets a callback that will be called when the upload progresses.
    ///
    /// The first argument contains the total size of the file, and the second
    /// the amount of bytes that have been uploaded.
    pub fn with_callback( mut self, callback: fn(usize, usize) ) -> Self {
        self.callback = Some(callback);

        self
    }

    /// Uploads a file chunk.
    fn upload_chunk( &self, chunk: &[u8] ) -> Result<Option<objects::File>, Error> {
        let chunk_length = chunk.len();

        let range_header = format!(
            "bytes {}-{}/{}",
            self.current_byte,
            self.current_byte + chunk_length - 1,
            self.total_bytes,
        );

        let request = Client::new().put(&self.upload_uri)
            .header( "Content-Length", chunk_length.to_string() )
            .header( "Content-Range",  range_header             )
            .body( chunk.to_vec() );

        let response = request.send()?;

        match response.status().as_u16() {
            200 | 202 => Ok( Some(serde_json::from_str( &response.text()? )?) ),
            #[cfg(not(tarpaulin_include))]
            308 => {
                // This should be received when and upload requires to be
                // resumed, however it seems to be received even in new uploads

                Ok(None)
            },
            #[cfg(not(tarpaulin_include))]
            _ => Err( response.into() ),
        }
    }

    pub fn upload_file( &mut self, file: &mut fs::File ) -> Result<objects::File, Error> {
        let size_of_file = file.metadata()?.len() as usize;

        self.current_byte = 0;
        self.total_bytes = size_of_file;

        if let Some(callback) = self.callback {
            (callback)(self.total_bytes, self.current_byte);
        }

        #[cfg(not(tarpaulin_include))]
        loop {
            let mut chunk = Vec::with_capacity(self.chunk_size);
            let read_bytes = file.by_ref().take(self.chunk_size as u64).read_to_end(&mut chunk)?;

            let uploaded_file = self.upload_chunk(&chunk)?;
            self.current_byte += read_bytes;

            if let Some(callback) = self.callback {
                (callback)(self.total_bytes, self.current_byte);
            }

            if let Some(file) = uploaded_file {
                return Ok(file)
            }

            if read_bytes == 0 || read_bytes < self.chunk_size {
                return Err( Error {
                    kind: ErrorKind::Request,
                    message: String::from("finished reading the file but the upload did not complete")
                } )
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::Write;
    use crate::ErrorKind;
    use super::{FileUploader, UPLOAD_CHUNK_SIZE};

    #[test]
    fn from_uri_test() {
        let upload_uri = String::from("resumable_session_uri");
        let uploader = FileUploader::from_uri(&upload_uri);

        assert_eq!(uploader.upload_uri, upload_uri);
        assert_eq!(uploader.current_byte, 0);
        assert_eq!(uploader.total_bytes, 0);
        assert_eq!(uploader.chunk_size, UPLOAD_CHUNK_SIZE);
        assert_eq!(uploader.callback, None);
    }

    #[test]
    fn with_callback_test() {
        fn callback( a: usize, b: usize ) {
            assert_eq!(a, 0);
            assert_eq!(b, 1);
        }

        let upload_uri = String::from("resumable_session_uri");
        let uploader = FileUploader::from_uri(&upload_uri)
            .with_callback(callback);

        assert!( uploader.callback.is_some() );

        (uploader.callback.unwrap())(0,1);
    }

    #[test]
    fn upload_test() {
        fn callback( total: usize, done: usize ) {
            assert_eq!(total, UPLOAD_CHUNK_SIZE * 4);
            assert_eq!(done, 0);
        }

        let mut uploader = FileUploader::from_uri("resumable_session_uri")
            .with_callback(callback);

        let empty_file_bytes: Vec<u8> = vec![0; UPLOAD_CHUNK_SIZE * 4];
        let mut test_file = testfile::create( |f| f.write_all(&empty_file_bytes) );

        let response = uploader.upload_file(&mut test_file);

        assert!( response.is_err() );
        assert_eq!( response.unwrap_err().kind, ErrorKind::Request );
    }
}