better_auth_core/store/schema/
check.rs

1use super::{SchemaFinding, SchemaMismatch};
2use crate::{AuthError, AuthResult};
3use async_trait::async_trait;
4use futures_util::{
5    FutureExt,
6    future::{BoxFuture, Shared},
7};
8use std::sync::Arc;
9use tokio::sync::Mutex;
10
11tokio::task_local! {
12    static TRANSACTIONS: Vec<Arc<SchemaCheck>>;
13}
14
15/// Adapter-owned introspection. A revision changes only after an explicit schema invalidation.
16#[async_trait]
17pub trait SchemaInspector: Send + Sync + 'static {
18    fn revision(&self) -> u64 {
19        0
20    }
21    async fn findings(&self) -> AuthResult<Vec<SchemaFinding>>;
22}
23
24#[derive(Debug, thiserror::Error)]
25pub enum SchemaCheckError {
26    #[error(transparent)]
27    Mismatch(#[from] SchemaMismatch),
28    /// Keep the original driver error while sharing one failed attempt with concurrent callers.
29    #[error(transparent)]
30    Store(#[from] AuthError),
31}
32
33type Verdict = Shared<BoxFuture<'static, Result<(), Arc<SchemaCheckError>>>>;
34
35#[derive(Default)]
36struct State {
37    revision: u64,
38    verdict: Option<Verdict>,
39}
40
41/// One cached verdict per configured adapter. Construct a new check for each auth build.
42pub struct SchemaCheck {
43    inspector: Arc<dyn SchemaInspector>,
44    state: Mutex<State>,
45}
46
47impl SchemaCheck {
48    pub fn new(inspector: Arc<dyn SchemaInspector>) -> Self {
49        Self {
50            inspector,
51            state: Mutex::new(State::default()),
52        }
53    }
54
55    pub async fn check(&self) -> AuthResult<()> {
56        loop {
57            let revision = self.inspector.revision();
58            let verdict = {
59                let mut state = self.state.lock().await;
60                if state.revision != revision {
61                    state.revision = revision;
62                    state.verdict = None;
63                }
64                state
65                    .verdict
66                    .get_or_insert_with(|| {
67                        let inspector = self.inspector.clone();
68                        async move {
69                            let findings = inspector
70                                .findings()
71                                .await
72                                .map_err(|error| Arc::new(SchemaCheckError::Store(error)))?;
73                            if findings.is_empty() {
74                                Ok(())
75                            } else {
76                                Err(Arc::new(SchemaCheckError::Mismatch(SchemaMismatch::new(
77                                    findings,
78                                ))))
79                            }
80                        }
81                        .boxed()
82                        .shared()
83                    })
84                    .clone()
85            };
86            let result = verdict.clone().await;
87            if self.inspector.revision() != revision {
88                continue;
89            }
90            if matches!(&result, Err(error) if matches!(error.as_ref(), SchemaCheckError::Store(_)))
91            {
92                let mut state = self.state.lock().await;
93                if state.revision == revision
94                    && state
95                        .verdict
96                        .as_ref()
97                        .is_some_and(|current| current.ptr_eq(&verdict))
98                {
99                    state.verdict = None;
100                }
101            }
102            return result.map_err(AuthError::SchemaCheck);
103        }
104    }
105}
106
107/// Explicit checks remain available when automatic runtime validation is disabled.
108#[derive(Clone)]
109pub struct SchemaValidation {
110    pub check: Arc<SchemaCheck>,
111    pub runtime_enabled: bool,
112}
113
114impl SchemaValidation {
115    pub async fn check_runtime(&self) -> AuthResult<()> {
116        let active = TRANSACTIONS
117            .try_with(|checks| checks.iter().any(|check| Arc::ptr_eq(check, &self.check)))
118            .unwrap_or(false);
119        if self.runtime_enabled && !active {
120            self.check.check().await?;
121        }
122        Ok(())
123    }
124
125    /// Scope the active transaction without treating another auth instance's check as completed.
126    pub async fn in_transaction<T>(&self, operation: impl std::future::Future<Output = T>) -> T {
127        let mut checks = TRANSACTIONS.try_with(Clone::clone).unwrap_or_default();
128        checks.push(self.check.clone());
129        TRANSACTIONS.scope(checks, Box::pin(operation)).await
130    }
131
132    /// Initialization reports failure but must not prevent the auth context from resolving.
133    pub fn start(&self) {
134        if !self.runtime_enabled {
135            return;
136        }
137        let check = self.check.clone();
138        let _task = tokio::spawn(async move {
139            if let Err(error) = check.check().await {
140                match &error {
141                    AuthError::SchemaCheck(error)
142                        if matches!(error.as_ref(), SchemaCheckError::Mismatch(_)) =>
143                    {
144                        tracing::error!("{error}")
145                    }
146                    _ => tracing::error!(
147                        "Could not validate the database schema. Check your database connection."
148                    ),
149                }
150            }
151        });
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
159    use tokio::sync::Notify;
160
161    struct Pending {
162        revision: AtomicU64,
163        calls: AtomicUsize,
164        started: Notify,
165        release: Notify,
166    }
167
168    #[async_trait]
169    impl SchemaInspector for Pending {
170        fn revision(&self) -> u64 {
171            self.revision.load(Ordering::Acquire)
172        }
173        async fn findings(&self) -> AuthResult<Vec<SchemaFinding>> {
174            if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
175                self.started.notify_one();
176                self.release.notified().await;
177                Ok(vec![SchemaFinding::MissingTable {
178                    table: "old_schema".into(),
179                }])
180            } else {
181                Ok(Vec::new())
182            }
183        }
184    }
185
186    #[tokio::test]
187    async fn invalidating_a_pending_check_discards_the_stale_mismatch_for_its_waiters() {
188        let inspector = Arc::new(Pending {
189            revision: AtomicU64::new(0),
190            calls: AtomicUsize::new(0),
191            started: Notify::new(),
192            release: Notify::new(),
193        });
194        let check = Arc::new(SchemaCheck::new(inspector.clone()));
195        let waiting = check.clone();
196        let task = tokio::spawn(async move { waiting.check().await });
197        inspector.started.notified().await;
198        inspector.revision.fetch_add(1, Ordering::AcqRel);
199        inspector.release.notify_one();
200        task.await.unwrap().unwrap();
201        check.check().await.unwrap();
202        assert_eq!(inspector.calls.load(Ordering::SeqCst), 2);
203    }
204
205    struct Retry(AtomicUsize);
206    #[async_trait]
207    impl SchemaInspector for Retry {
208        async fn findings(&self) -> AuthResult<Vec<SchemaFinding>> {
209            if self.0.fetch_add(1, Ordering::SeqCst) == 0 {
210                Err(AuthError::Database(
211                    crate::error::DatabaseError::Connection("closed fixture handle".into()),
212                ))
213            } else {
214                Ok(Vec::new())
215            }
216        }
217    }
218
219    #[tokio::test]
220    async fn a_driver_failure_can_retry_and_then_caches_the_clean_result() {
221        let inspector = Arc::new(Retry(AtomicUsize::new(0)));
222        let check = SchemaCheck::new(inspector.clone());
223        let failure = check.check().await.unwrap_err();
224        assert!(
225            matches!(failure, AuthError::SchemaCheck(ref error) if matches!(error.as_ref(), SchemaCheckError::Store(AuthError::Database(_))))
226        );
227        check.check().await.unwrap();
228        check.check().await.unwrap();
229        assert_eq!(inspector.0.load(Ordering::SeqCst), 2);
230    }
231}
better_auth_core/store/schema/check.rs

better_auth_core/store/schema/
check.rs
