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

use crate::message::Response;
use crate::pool::GremlinConnectionManager;
use crate::structure::GValue;
use crate::{GremlinClient, GremlinResult};
use r2d2::PooledConnection;
use std::collections::VecDeque;

#[derive(Debug)]
pub struct GResultSet {
    client: GremlinClient,
    results: VecDeque<GValue>,
    response: Response,
    conn: PooledConnection<GremlinConnectionManager>,
}

impl GResultSet {
    pub(crate) fn new(
        client: GremlinClient,
        results: VecDeque<GValue>,
        response: Response,
        conn: PooledConnection<GremlinConnectionManager>,
    ) -> GResultSet {
        GResultSet {
            client,
            results,
            response,
            conn,
        }
    }

    fn fetch(&mut self) -> GremlinResult<()> {
        if self.results.is_empty() && self.response.status.code == 206 {
            let (response, resuts) = self.client.read_response(&mut self.conn)?;
            self.response = response;
            self.results = resuts;
        }
        Ok(())
    }
}

impl Iterator for GResultSet {
    type Item = GremlinResult<GValue>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.fetch() {
            Ok(_) => self.results.pop_front().map(Ok),
            Err(e) => Some(Err(e)),
        }
    }
}