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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use crate::{
    config::Database,
    errors::{unexpected, Result},
    messages::{BoltRequest, BoltResponse},
    pool::ManagedConnection,
    query::Query,
    stream::RowStream,
};

/// A handle which is used to control a transaction, created as a result of [`crate::Graph::start_txn`]
///
/// When a transation is started, a dedicated connection is resered and moved into the handle which
/// will be released to the connection pool when the [`Txn`] handle is dropped.
pub struct Txn {
    db: Database,
    fetch_size: usize,
    connection: ManagedConnection,
}

impl Txn {
    pub(crate) async fn new(
        db: Database,
        fetch_size: usize,
        mut connection: ManagedConnection,
    ) -> Result<Self> {
        let begin = BoltRequest::begin(&db);
        match connection.send_recv(begin).await? {
            BoltResponse::Success(_) => Ok(Txn {
                db,
                fetch_size,
                connection,
            }),
            msg => Err(unexpected(msg, "BEGIN")),
        }
    }

    /// Runs multiple queries one after the other in the same connection
    pub async fn run_queries<Q: Into<Query>>(
        &mut self,
        queries: impl IntoIterator<Item = Q>,
    ) -> Result<()> {
        for query in queries {
            self.run(query.into()).await?;
        }
        Ok(())
    }

    /// Runs a single query and discards the stream.
    pub async fn run(&mut self, q: Query) -> Result<()> {
        q.run(&self.db, &mut self.connection).await
    }

    /// Executes a query and returns a [`RowStream`]
    pub async fn execute(&mut self, q: Query) -> Result<RowStream> {
        q.execute_mut(&self.db, self.fetch_size, &mut self.connection)
            .await
    }

    /// Commits the transaction in progress
    pub async fn commit(mut self) -> Result<()> {
        let commit = BoltRequest::commit();
        match self.connection.send_recv(commit).await? {
            BoltResponse::Success(_) => Ok(()),
            msg => Err(unexpected(msg, "COMMIT")),
        }
    }

    /// rollback/abort the current transaction
    pub async fn rollback(mut self) -> Result<()> {
        let rollback = BoltRequest::rollback();
        match self.connection.send_recv(rollback).await? {
            BoltResponse::Success(_) => Ok(()),
            msg => Err(unexpected(msg, "ROLLBACK")),
        }
    }

    pub fn handle(&mut self) -> &mut impl TransactionHandle {
        self
    }
}

const _: () = {
    const fn assert_send_sync<T: ?Sized + Send + Sync>() {}
    assert_send_sync::<Txn>();
};

pub trait TransactionHandle: private::Handle {}

impl TransactionHandle for Txn {}
impl TransactionHandle for ManagedConnection {}
impl<T: TransactionHandle> TransactionHandle for &mut T {}

pub(crate) mod private {
    use crate::{pool::ManagedConnection, Txn};

    pub trait Handle {
        fn connection(&mut self) -> &mut ManagedConnection;
    }

    impl Handle for Txn {
        fn connection(&mut self) -> &mut ManagedConnection {
            &mut self.connection
        }
    }

    impl Handle for ManagedConnection {
        fn connection(&mut self) -> &mut ManagedConnection {
            self
        }
    }

    impl<T: Handle> Handle for &mut T {
        fn connection(&mut self) -> &mut ManagedConnection {
            (**self).connection()
        }
    }
}