torii_storage_postgres/
session.rs

1use crate::PostgresStorage;
2use async_trait::async_trait;
3use chrono::{DateTime, Utc};
4use torii_core::error::StorageError;
5use torii_core::session::SessionToken;
6use torii_core::{Session, SessionStorage, UserId};
7
8#[derive(Debug, Clone, sqlx::FromRow)]
9pub struct PostgresSession {
10    pub id: Option<i64>,
11    pub user_id: String,
12    pub token: String,
13    pub user_agent: Option<String>,
14    pub ip_address: Option<String>,
15    pub created_at: DateTime<Utc>,
16    pub updated_at: DateTime<Utc>,
17    pub expires_at: DateTime<Utc>,
18}
19
20impl From<PostgresSession> for Session {
21    fn from(session: PostgresSession) -> Self {
22        Session::builder()
23            .token(SessionToken::new(&session.token))
24            .user_id(UserId::new(&session.user_id))
25            .user_agent(session.user_agent)
26            .ip_address(session.ip_address)
27            .created_at(session.created_at)
28            .updated_at(session.updated_at)
29            .expires_at(session.expires_at)
30            .build()
31            .unwrap()
32    }
33}
34
35impl From<Session> for PostgresSession {
36    fn from(session: Session) -> Self {
37        PostgresSession {
38            id: None,
39            token: session.token.into_inner(),
40            user_id: session.user_id.into_inner(),
41            user_agent: session.user_agent,
42            ip_address: session.ip_address,
43            created_at: session.created_at,
44            updated_at: session.updated_at,
45            expires_at: session.expires_at,
46        }
47    }
48}
49
50#[async_trait]
51impl SessionStorage for PostgresStorage {
52    type Error = torii_core::Error;
53
54    async fn create_session(&self, session: &Session) -> Result<Session, Self::Error> {
55        sqlx::query("INSERT INTO sessions (token, user_id, user_agent, ip_address, created_at, updated_at, expires_at) VALUES ($1, $2, $3, $4, $5, $6, $7)")
56            .bind(session.token.as_str())
57            .bind(session.user_id.as_str())
58            .bind(&session.user_agent)
59            .bind(&session.ip_address)
60            .bind(session.created_at)
61            .bind(session.updated_at)
62            .bind(session.expires_at)
63            .execute(&self.pool)
64            .await
65            .map_err(|e| {
66                tracing::error!(error = %e, "Failed to create session");
67                StorageError::Database("Failed to create session".to_string())
68            })?;
69
70        Ok(self.get_session(&session.token).await?.unwrap())
71    }
72
73    async fn get_session(&self, token: &SessionToken) -> Result<Option<Session>, Self::Error> {
74        let session = sqlx::query_as::<_, PostgresSession>(
75            r#"
76            SELECT id, token, user_id, user_agent, ip_address, created_at, updated_at, expires_at
77            FROM sessions
78            WHERE token = $1
79            "#,
80        )
81        .bind(token.as_str())
82        .fetch_one(&self.pool)
83        .await
84        .map_err(|e| {
85            tracing::error!(error = %e, "Failed to get session");
86            StorageError::Database("Failed to get session".to_string())
87        })?;
88
89        Ok(Some(session.into()))
90    }
91
92    async fn delete_session(&self, token: &SessionToken) -> Result<(), Self::Error> {
93        sqlx::query("DELETE FROM sessions WHERE token = $1")
94            .bind(token.as_str())
95            .execute(&self.pool)
96            .await
97            .map_err(|e| {
98                tracing::error!(error = %e, "Failed to delete session");
99                StorageError::Database("Failed to delete session".to_string())
100            })?;
101
102        Ok(())
103    }
104
105    async fn cleanup_expired_sessions(&self) -> Result<(), Self::Error> {
106        sqlx::query("DELETE FROM sessions WHERE expires_at < $1")
107            .bind(Utc::now())
108            .execute(&self.pool)
109            .await
110            .map_err(|e| {
111                tracing::error!(error = %e, "Failed to cleanup expired sessions");
112                StorageError::Database("Failed to cleanup expired sessions".to_string())
113            })?;
114
115        Ok(())
116    }
117
118    async fn delete_sessions_for_user(&self, user_id: &UserId) -> Result<(), Self::Error> {
119        sqlx::query("DELETE FROM sessions WHERE user_id = $1")
120            .bind(user_id.as_str())
121            .execute(&self.pool)
122            .await
123            .map_err(|e| {
124                tracing::error!(error = %e, "Failed to delete sessions for user");
125                StorageError::Database("Failed to delete sessions for user".to_string())
126            })?;
127
128        Ok(())
129    }
130}
131
132#[cfg(test)]
133mod test {
134    use super::*;
135    use std::time::Duration;
136    use torii_core::session::SessionToken;
137    use torii_core::{Session, UserId, UserStorage};
138
139    #[tokio::test]
140    async fn test_postgres_storage() {
141        let storage = crate::tests::setup_test_db().await;
142        let user_id = UserId::new_random();
143        let user = crate::tests::create_test_user(&storage, &user_id)
144            .await
145            .expect("Failed to create user");
146        assert_eq!(user.email, format!("test{user_id}@example.com"));
147
148        let fetched = storage
149            .get_user(&user_id)
150            .await
151            .expect("Failed to get user");
152        assert_eq!(
153            fetched.expect("User should exist").email,
154            format!("test{user_id}@example.com")
155        );
156
157        storage
158            .delete_user(&user_id)
159            .await
160            .expect("Failed to delete user");
161        let deleted = storage
162            .get_user(&user_id)
163            .await
164            .expect("Failed to get user");
165        assert!(deleted.is_none());
166    }
167
168    #[tokio::test]
169    async fn test_postgres_session_storage() {
170        let storage = crate::tests::setup_test_db().await;
171        let user_id = UserId::new_random();
172        let session_id = SessionToken::new_random();
173        crate::tests::create_test_user(&storage, &user_id)
174            .await
175            .expect("Failed to create user");
176
177        let _session = crate::tests::create_test_session(
178            &storage,
179            &session_id,
180            &user_id,
181            Duration::from_secs(1000),
182        )
183        .await
184        .expect("Failed to create session");
185
186        let fetched = storage
187            .get_session(&session_id)
188            .await
189            .expect("Failed to get session");
190        assert_eq!(fetched.unwrap().user_id, user_id);
191
192        storage
193            .delete_session(&session_id)
194            .await
195            .expect("Failed to delete session");
196        let deleted = storage.get_session(&session_id).await;
197        assert!(deleted.is_err());
198    }
199
200    #[tokio::test]
201    async fn test_postgres_session_cleanup() {
202        let storage = crate::tests::setup_test_db().await;
203        let user_id = UserId::new_random();
204        crate::tests::create_test_user(&storage, &user_id)
205            .await
206            .expect("Failed to create user");
207
208        // Create an already expired session by setting expires_at in the past
209        let session_id = SessionToken::new_random();
210        let expired_session = Session {
211            token: SessionToken::new_random(),
212            user_id: user_id.clone(),
213            user_agent: None,
214            ip_address: None,
215            created_at: chrono::Utc::now(),
216            updated_at: chrono::Utc::now(),
217            expires_at: chrono::Utc::now() - chrono::Duration::seconds(1),
218        };
219        storage
220            .create_session(&expired_session)
221            .await
222            .expect("Failed to create expired session");
223
224        // Create valid session
225        crate::tests::create_test_session(
226            &storage,
227            &session_id,
228            &user_id,
229            Duration::from_secs(3600),
230        )
231        .await
232        .expect("Failed to create valid session");
233
234        // Run cleanup
235        storage
236            .cleanup_expired_sessions()
237            .await
238            .expect("Failed to cleanup sessions");
239
240        // Verify expired session was removed
241        let expired_session = storage.get_session(&expired_session.token).await;
242        assert!(expired_session.is_err());
243
244        // Verify valid session remains
245        let valid_session = storage
246            .get_session(&session_id)
247            .await
248            .expect("Failed to get valid session");
249        assert_eq!(valid_session.unwrap().user_id, user_id);
250    }
251
252    #[tokio::test]
253    async fn test_delete_sessions_for_user() {
254        let storage = crate::tests::setup_test_db().await;
255
256        // Create test users
257        let user_id1 = UserId::new_random();
258        crate::tests::create_test_user(&storage, &user_id1)
259            .await
260            .expect("Failed to create user 1");
261        let user_id2 = UserId::new_random();
262        crate::tests::create_test_user(&storage, &user_id2)
263            .await
264            .expect("Failed to create user 2");
265
266        // Create sessions for user 1
267        let session_id1 = SessionToken::new_random();
268        crate::tests::create_test_session(
269            &storage,
270            &session_id1,
271            &user_id1,
272            Duration::from_secs(3600),
273        )
274        .await
275        .expect("Failed to create session 1");
276        let session_id2 = SessionToken::new_random();
277        crate::tests::create_test_session(
278            &storage,
279            &session_id2,
280            &user_id1,
281            Duration::from_secs(3600),
282        )
283        .await
284        .expect("Failed to create session 2");
285
286        // Create session for user 2
287        let session_id3 = SessionToken::new_random();
288        crate::tests::create_test_session(
289            &storage,
290            &session_id3,
291            &user_id2,
292            Duration::from_secs(3600),
293        )
294        .await
295        .expect("Failed to create session 3");
296
297        // Delete all sessions for user 1
298        storage
299            .delete_sessions_for_user(&user_id1)
300            .await
301            .expect("Failed to delete sessions for user");
302
303        // Verify user 1's sessions are deleted
304        let session1 = storage.get_session(&session_id1).await;
305        assert!(session1.is_err());
306        let session2 = storage.get_session(&session_id2).await;
307        assert!(session2.is_err());
308
309        // Verify user 2's session remains
310        let session3 = storage
311            .get_session(&session_id3)
312            .await
313            .expect("Failed to get session 3");
314        assert_eq!(session3.unwrap().user_id, user_id2);
315    }
316}
torii_storage_postgres/session.rs

torii_storage_postgres/
session.rs
