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

torii_storage_postgres/
session.rs
