grafbase_sdk/host_io/postgres/
connection.rs

1use std::fmt;
2
3use crate::{SdkError, wit};
4
5/// A Postgres database connection.
6///
7/// This represents a single connection to the Postgres database,
8/// which can be used to execute queries or perform other database operations.
9pub struct Connection(pub(crate) wit::PgConnection);
10
11impl fmt::Debug for Connection {
12    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
13        f.write_str("Connection { ... }")
14    }
15}
16
17impl From<wit::PgConnection> for Connection {
18    fn from(conn: wit::PgConnection) -> Self {
19        Self(conn)
20    }
21}
22
23/// A Postgres database transaction.
24///
25/// This represents an active database transaction that can be either committed
26/// to make changes permanent, or rolled back to discard changes.
27///
28/// The transaction must be either commited or rolled back, otherwise it will be
29/// automatically rolled back when dropped.
30pub struct Transaction {
31    pub(crate) inner: wit::PgTransaction,
32    committed_or_rolled_back: bool,
33}
34
35impl From<wit::PgTransaction> for Transaction {
36    fn from(inner: wit::PgTransaction) -> Self {
37        Self {
38            inner,
39            committed_or_rolled_back: false,
40        }
41    }
42}
43
44impl fmt::Debug for Transaction {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        f.write_str("Transaction { ... }")
47    }
48}
49
50impl Transaction {
51    /// Commits the transaction, making all changes permanent.
52    ///
53    /// # Returns
54    /// `Ok(())` if the transaction was successfully committed, or an error message
55    pub fn commit(mut self) -> Result<(), SdkError> {
56        self.committed_or_rolled_back = true;
57
58        self.inner.commit().map_err(SdkError::from)
59    }
60
61    /// Rolls back the transaction, discarding all changes.
62    ///
63    /// # Returns
64    /// `Ok(())` if the transaction was successfully rolled back, or an error message
65    pub fn rollback(mut self) -> Result<(), SdkError> {
66        self.committed_or_rolled_back = true;
67
68        self.inner.rollback().map_err(SdkError::from)
69    }
70}
71
72impl Drop for Transaction {
73    fn drop(&mut self) {
74        if !self.committed_or_rolled_back {
75            self.inner.rollback().unwrap()
76        }
77    }
78}
79
80/// Represents either a [`Connection`] or a [`Transaction`].
81///
82/// This enum is used in functions that can operate on either a regular
83/// connection or within an active transaction, allowing for flexibility
84/// in how database operations are performed.
85#[derive(Clone, Copy, Debug)]
86pub enum ConnectionLike<'a> {
87    /// A regular database connection, not part of an explicit transaction.
88    Connection(&'a Connection),
89    /// An active database transaction. Operations performed using this variant
90    /// will be part of the ongoing transaction.
91    Transaction(&'a Transaction),
92}
93
94impl<'a> From<&'a Connection> for ConnectionLike<'a> {
95    fn from(connection: &'a Connection) -> Self {
96        ConnectionLike::Connection(connection)
97    }
98}
99
100impl<'a> From<&'a Transaction> for ConnectionLike<'a> {
101    fn from(transaction: &'a Transaction) -> Self {
102        ConnectionLike::Transaction(transaction)
103    }
104}
105
106impl ConnectionLike<'_> {
107    pub(crate) fn query<'a>(
108        &'a self,
109        query: &'a str,
110        params: (&[wit::PgBoundValue], &[wit::PgValue]),
111    ) -> Result<Vec<wit::PgRow>, SdkError> {
112        match self {
113            ConnectionLike::Connection(connection) => connection.0.query(query, params).map_err(SdkError::from),
114            ConnectionLike::Transaction(transaction) => transaction.inner.query(query, params).map_err(SdkError::from),
115        }
116    }
117
118    pub(crate) fn execute<'a>(
119        &'a self,
120        query: &'a str,
121        params: (&[wit::PgBoundValue], &[wit::PgValue]),
122    ) -> Result<u64, SdkError> {
123        match self {
124            ConnectionLike::Connection(connection) => connection.0.execute(query, params).map_err(SdkError::from),
125            ConnectionLike::Transaction(transaction) => {
126                transaction.inner.execute(query, params).map_err(SdkError::from)
127            }
128        }
129    }
130}
grafbase_sdk/host_io/postgres/connection.rs

grafbase_sdk/host_io/postgres/
connection.rs
