systemprompt_analytics/repository/session/
mod.rs

1//! `SessionRepository` — repository surface over `user_sessions`. Internal
2//! submodules: `queries` (read), `mutations` (write), `behavioral` (bot
3//! detection writes), `behavioral_queries` (bot detection reads), `types`
4//! (DTO/row structs).
5//!
6//! Copyright (c) systemprompt.io — Business Source License 1.1.
7//! See <https://systemprompt.io> for licensing details.
8
9mod behavioral;
10mod behavioral_queries;
11mod mutations;
12mod queries;
13mod types;
14
15use std::sync::Arc;
16
17use crate::Result;
18use chrono::{DateTime, Utc};
19use sqlx::PgPool;
20use systemprompt_database::DbPool;
21use systemprompt_identifiers::{SessionId, UserId};
22
23use crate::models::AnalyticsSession;
24
25pub(super) use types::ActiveSessionLookup;
26pub use types::{
27    CreateSessionParams, SessionBehavioralData, SessionMigrationResult, SessionRecord,
28};
29
30#[derive(Clone, Debug)]
31pub struct SessionRepository {
32    pool: Arc<PgPool>,
33    write_pool: Arc<PgPool>,
34}
35
36impl SessionRepository {
37    pub fn new(db: &DbPool) -> Result<Self> {
38        let pool = db.pool_arc()?;
39        let write_pool = db.write_pool_arc()?;
40        Ok(Self { pool, write_pool })
41    }
42
43    pub async fn find_by_id(&self, session_id: &SessionId) -> Result<Option<AnalyticsSession>> {
44        queries::find_by_id(&self.pool, session_id).await
45    }
46
47    pub async fn find_active_by_id(
48        &self,
49        session_id: &SessionId,
50    ) -> Result<Option<ActiveSessionLookup>> {
51        queries::find_active_by_id(&self.pool, session_id).await
52    }
53
54    pub async fn revoke_session(&self, session_id: &SessionId) -> Result<()> {
55        mutations::revoke_session(&self.write_pool, session_id).await
56    }
57
58    pub async fn revoke_all_for_user(&self, user_id: &UserId) -> Result<u64> {
59        mutations::revoke_all_for_user(&self.write_pool, user_id).await
60    }
61
62    pub async fn find_by_fingerprint(
63        &self,
64        fingerprint_hash: &str,
65        user_id: &UserId,
66    ) -> Result<Option<AnalyticsSession>> {
67        queries::find_by_fingerprint(&self.pool, fingerprint_hash, user_id).await
68    }
69
70    pub async fn list_active_by_user(&self, user_id: &UserId) -> Result<Vec<AnalyticsSession>> {
71        queries::list_active_by_user(&self.pool, user_id).await
72    }
73
74    pub async fn update_activity(&self, session_id: &SessionId) -> Result<()> {
75        mutations::update_activity(&self.write_pool, session_id).await
76    }
77
78    pub async fn increment_request_count(&self, session_id: &SessionId) -> Result<()> {
79        mutations::increment_request_count(&self.write_pool, session_id).await
80    }
81
82    pub async fn increment_task_count(&self, session_id: &SessionId) -> Result<()> {
83        mutations::increment_task_count(&self.write_pool, session_id).await
84    }
85
86    pub async fn increment_message_count(&self, session_id: &SessionId) -> Result<()> {
87        mutations::increment_message_count(&self.write_pool, session_id).await
88    }
89
90    pub async fn end_session(&self, session_id: &SessionId) -> Result<()> {
91        mutations::end_session(&self.write_pool, session_id).await
92    }
93
94    pub async fn mark_as_scanner(&self, session_id: &SessionId) -> Result<()> {
95        mutations::mark_as_scanner(&self.write_pool, session_id).await
96    }
97
98    pub async fn mark_converted(&self, session_id: &SessionId) -> Result<()> {
99        mutations::mark_converted(&self.write_pool, session_id).await
100    }
101
102    pub async fn mark_as_behavioral_bot(&self, session_id: &SessionId, reason: &str) -> Result<()> {
103        behavioral::mark_as_behavioral_bot(&self.write_pool, session_id, reason).await
104    }
105
106    pub async fn check_and_mark_behavioral_bot(
107        &self,
108        session_id: &SessionId,
109        request_count_threshold: i32,
110    ) -> Result<bool> {
111        behavioral::check_and_mark_behavioral_bot(
112            &self.write_pool,
113            session_id,
114            request_count_threshold,
115        )
116        .await
117    }
118
119    pub async fn cleanup_inactive(&self, inactive_hours: i32) -> Result<u64> {
120        mutations::cleanup_inactive(&self.write_pool, inactive_hours).await
121    }
122
123    pub async fn count_inactive(&self, inactive_hours: i32) -> Result<i64> {
124        queries::count_inactive(&self.pool, inactive_hours).await
125    }
126
127    pub async fn backfill_session_geo(
128        &self,
129        geoip_reader: Option<&crate::GeoIpReader>,
130        batch_size: i64,
131    ) -> Result<u64> {
132        mutations::backfill_session_geo(&self.write_pool, geoip_reader, batch_size).await
133    }
134
135    pub async fn count_sessions_missing_geo(&self) -> Result<i64> {
136        mutations::count_sessions_missing_geo(&self.pool).await
137    }
138
139    pub async fn migrate_user_sessions(
140        &self,
141        old_user_id: &UserId,
142        new_user_id: &UserId,
143    ) -> Result<u64> {
144        mutations::migrate_user_sessions(&self.write_pool, old_user_id, new_user_id).await
145    }
146
147    pub async fn create_session(&self, params: &CreateSessionParams<'_>) -> Result<()> {
148        mutations::create_session(&self.write_pool, params).await
149    }
150
151    pub async fn find_recent_by_fingerprint(
152        &self,
153        fingerprint_hash: &str,
154        max_age_seconds: i64,
155    ) -> Result<Option<SessionRecord>> {
156        queries::find_recent_by_fingerprint(&self.pool, fingerprint_hash, max_age_seconds).await
157    }
158
159    pub async fn increment_ai_usage(
160        &self,
161        session_id: &SessionId,
162        tokens: i32,
163        cost_microdollars: i64,
164    ) -> Result<()> {
165        mutations::increment_ai_usage(&self.write_pool, session_id, tokens, cost_microdollars).await
166    }
167
168    pub async fn update_behavioral_detection(
169        &self,
170        session_id: &SessionId,
171        score: i32,
172        is_behavioral_bot: bool,
173        reason: Option<&str>,
174    ) -> Result<()> {
175        behavioral::update_behavioral_detection(
176            &self.write_pool,
177            session_id,
178            score,
179            is_behavioral_bot,
180            reason,
181        )
182        .await
183    }
184
185    pub async fn count_sessions_by_fingerprint(
186        &self,
187        fingerprint_hash: &str,
188        window_hours: i64,
189    ) -> Result<i64> {
190        behavioral_queries::count_sessions_by_fingerprint(
191            &self.pool,
192            fingerprint_hash,
193            window_hours,
194        )
195        .await
196    }
197
198    pub async fn get_endpoint_sequence(&self, session_id: &SessionId) -> Result<Vec<String>> {
199        behavioral_queries::get_endpoint_sequence(&self.pool, session_id).await
200    }
201
202    pub async fn get_request_timestamps(
203        &self,
204        session_id: &SessionId,
205    ) -> Result<Vec<DateTime<Utc>>> {
206        behavioral_queries::get_request_timestamps(&self.pool, session_id).await
207    }
208
209    pub async fn get_total_content_pages(&self) -> Result<i64> {
210        queries::get_total_content_pages(&self.pool).await
211    }
212
213    pub async fn get_session_for_behavioral_analysis(
214        &self,
215        session_id: &SessionId,
216    ) -> Result<Option<SessionBehavioralData>> {
217        behavioral_queries::get_session_for_behavioral_analysis(&self.pool, session_id).await
218    }
219
220    pub async fn has_analytics_events(&self, session_id: &SessionId) -> Result<bool> {
221        behavioral_queries::has_analytics_events(&self.pool, session_id).await
222    }
223
224    pub async fn count_unique_ips_by_fingerprint(
225        &self,
226        fingerprint_hash: &str,
227        window_days: i64,
228    ) -> Result<i64> {
229        behavioral_queries::count_unique_ips_by_fingerprint(
230            &self.pool,
231            fingerprint_hash,
232            window_days,
233        )
234        .await
235    }
236
237    pub async fn count_engagement_events_by_fingerprint(
238        &self,
239        fingerprint_hash: &str,
240        window_days: i64,
241    ) -> Result<i64> {
242        behavioral_queries::count_engagement_events_by_fingerprint(
243            &self.pool,
244            fingerprint_hash,
245            window_days,
246        )
247        .await
248    }
249
250    pub async fn get_session_starts_by_fingerprint(
251        &self,
252        fingerprint_hash: &str,
253        window_days: i64,
254    ) -> Result<Vec<DateTime<Utc>>> {
255        behavioral_queries::get_session_starts_by_fingerprint(
256            &self.pool,
257            fingerprint_hash,
258            window_days,
259        )
260        .await
261    }
262
263    pub async fn get_session_velocity(
264        &self,
265        session_id: &SessionId,
266    ) -> Result<(Option<i64>, Option<i64>)> {
267        behavioral_queries::get_session_velocity(&self.pool, session_id).await
268    }
269}
systemprompt_analytics/repository/session/mod.rs

systemprompt_analytics/repository/session/
mod.rs
