closure_tree/
lock.rs

1use sea_orm::{
2    ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, Statement,
3    TransactionTrait, Value,
4};
5
6use crate::config::AdvisoryLockStrategy;
7use crate::error::ClosureTreeError;
8
9pub struct LockedTransaction {
10    txn: Option<DatabaseTransaction>,
11    key: Option<String>,
12}
13
14impl LockedTransaction {
15    pub async fn acquire(
16        strategy: &AdvisoryLockStrategy,
17        db: &DatabaseConnection,
18    ) -> Result<Self, ClosureTreeError> {
19        let key = match strategy {
20            AdvisoryLockStrategy::Disabled => None,
21            AdvisoryLockStrategy::Namespaced(key) => Some(key.as_str().to_owned()),
22        };
23
24        let txn = db.begin().await?;
25
26        if let Some(ref key) = key {
27            if let Err(err) = acquire_lock(&txn, key).await {
28                let _ = txn.rollback().await;
29                return Err(err);
30            }
31        }
32
33        Ok(Self {
34            txn: Some(txn),
35            key,
36        })
37    }
38
39    pub fn connection(&self) -> &DatabaseTransaction {
40        self.txn.as_ref().expect("transaction already consumed")
41    }
42
43    pub async fn commit(mut self) -> Result<(), ClosureTreeError> {
44        if let Some(ref key) = self.key {
45            if let Some(txn) = self.txn.as_ref() {
46                release_lock(txn, key).await?;
47            }
48        }
49
50        if let Some(txn) = self.txn.take() {
51            txn.commit().await?;
52        }
53
54        Ok(())
55    }
56
57    pub async fn rollback(mut self) -> Result<(), ClosureTreeError> {
58        if let Some(ref key) = self.key {
59            if let Some(txn) = self.txn.as_ref() {
60                let _ = release_lock(txn, key).await;
61            }
62        }
63
64        if let Some(txn) = self.txn.take() {
65            txn.rollback().await?;
66        }
67
68        Ok(())
69    }
70}
71
72async fn acquire_lock(txn: &DatabaseTransaction, key: &str) -> Result<(), ClosureTreeError> {
73    txn.execute(Statement::from_sql_and_values(
74        DbBackend::Postgres,
75        "SELECT pg_advisory_lock(hashtext($1), 0)",
76        vec![Value::from(key)],
77    ))
78    .await?;
79    Ok(())
80}
81
82async fn release_lock(txn: &DatabaseTransaction, key: &str) -> Result<(), ClosureTreeError> {
83    txn.execute(Statement::from_sql_and_values(
84        DbBackend::Postgres,
85        "SELECT pg_advisory_unlock(hashtext($1), 0)",
86        vec![Value::from(key)],
87    ))
88    .await?;
89    Ok(())
90}
closure_tree/lock.rs

closure_tree/
lock.rs
