1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

use fallible_iterator::FallibleIterator;
use rusoto_core::DispatchSignedRequest;
use rusoto_credential::ProvideAwsCredentials;
use rusoto_s3::{ListObjectsV2Error, ListObjectsV2Request, Object, S3, S3Client};
use std::vec::IntoIter;

/// Iterator over all objects or objects with a given prefix
pub struct ObjectIter<'a, P, D>
where
    P: 'a + ProvideAwsCredentials,
    D: 'a + DispatchSignedRequest,
{
    client: &'a S3Client<P, D>,
    request: ListObjectsV2Request,
    objects: IntoIter<Object>,
    exhausted: bool,
}

impl<'a, P, D> ObjectIter<'a, P, D>
where
    P: ProvideAwsCredentials,
    D: DispatchSignedRequest,
{
    pub(crate) fn new(client: &'a S3Client<P, D>, bucket: &str, prefix: Option<&str>) -> Self {
        let request = ListObjectsV2Request {
            bucket: bucket.to_owned(),
            max_keys: Some(1000),
            prefix: prefix.map(|s| s.to_owned()),
            ..Default::default()
        };

        ObjectIter {
            client,
            request,
            objects: Vec::new().into_iter(),
            exhausted: false,
        }
    }

    fn next_objects(&mut self) -> Result<(), ListObjectsV2Error> {
        let resp = self.client.list_objects_v2(&self.request)?;
        self.objects = resp.contents.unwrap_or_else(Vec::new).into_iter();
        match resp.next_continuation_token {
            next @ Some(_) => self.request.continuation_token = next,
            None => self.exhausted = true,
        };
        Ok(())
    }
}

impl<'a, P, D> FallibleIterator for ObjectIter<'a, P, D>
where
    P: ProvideAwsCredentials,
    D: DispatchSignedRequest,
{
    type Item = Object;
    type Error = ListObjectsV2Error;

    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
        if let object @ Some(_) = self.objects.next() {
            return Ok(object);
        }

        if !self.exhausted {
            self.next_objects()?;
            Ok(self.objects.next())
        } else {
            Ok(None)
        }
    }

    fn nth(&mut self, mut n: usize) -> Result<Option<Self::Item>, Self::Error> {
        while self.objects.len() <= n {
            if self.exhausted {
                return Ok(None);
            }
            n -= self.objects.len();
            self.next_objects()?;
        }
        Ok(self.objects.nth(n))
    }

    fn count(mut self) -> Result<usize, Self::Error> {
        let mut count = self.objects.len();
        while !self.exhausted {
            self.next_objects()?;
            count += self.objects.len();
        }
        Ok(count)
    }
}