systemprompt_analytics/repository/session/
mod.rs1mod 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 migrate_user_sessions(
128 &self,
129 old_user_id: &UserId,
130 new_user_id: &UserId,
131 ) -> Result<u64> {
132 mutations::migrate_user_sessions(&self.write_pool, old_user_id, new_user_id).await
133 }
134
135 pub async fn create_session(&self, params: &CreateSessionParams<'_>) -> Result<()> {
136 mutations::create_session(&self.write_pool, params).await
137 }
138
139 pub async fn find_recent_by_fingerprint(
140 &self,
141 fingerprint_hash: &str,
142 max_age_seconds: i64,
143 ) -> Result<Option<SessionRecord>> {
144 queries::find_recent_by_fingerprint(&self.pool, fingerprint_hash, max_age_seconds).await
145 }
146
147 pub async fn exists(&self, session_id: &SessionId) -> Result<bool> {
148 queries::exists(&self.pool, session_id).await
149 }
150
151 pub async fn increment_ai_usage(
152 &self,
153 session_id: &SessionId,
154 tokens: i32,
155 cost_microdollars: i64,
156 ) -> Result<()> {
157 mutations::increment_ai_usage(&self.write_pool, session_id, tokens, cost_microdollars).await
158 }
159
160 pub async fn update_behavioral_detection(
161 &self,
162 session_id: &SessionId,
163 score: i32,
164 is_behavioral_bot: bool,
165 reason: Option<&str>,
166 ) -> Result<()> {
167 behavioral::update_behavioral_detection(
168 &self.write_pool,
169 session_id,
170 score,
171 is_behavioral_bot,
172 reason,
173 )
174 .await
175 }
176
177 pub async fn count_sessions_by_fingerprint(
178 &self,
179 fingerprint_hash: &str,
180 window_hours: i64,
181 ) -> Result<i64> {
182 behavioral_queries::count_sessions_by_fingerprint(
183 &self.pool,
184 fingerprint_hash,
185 window_hours,
186 )
187 .await
188 }
189
190 pub async fn get_endpoint_sequence(&self, session_id: &SessionId) -> Result<Vec<String>> {
191 behavioral_queries::get_endpoint_sequence(&self.pool, session_id).await
192 }
193
194 pub async fn get_request_timestamps(
195 &self,
196 session_id: &SessionId,
197 ) -> Result<Vec<DateTime<Utc>>> {
198 behavioral_queries::get_request_timestamps(&self.pool, session_id).await
199 }
200
201 pub async fn get_total_content_pages(&self) -> Result<i64> {
202 queries::get_total_content_pages(&self.pool).await
203 }
204
205 pub async fn get_session_for_behavioral_analysis(
206 &self,
207 session_id: &SessionId,
208 ) -> Result<Option<SessionBehavioralData>> {
209 behavioral_queries::get_session_for_behavioral_analysis(&self.pool, session_id).await
210 }
211
212 pub async fn has_analytics_events(&self, session_id: &SessionId) -> Result<bool> {
213 behavioral_queries::has_analytics_events(&self.pool, session_id).await
214 }
215
216 pub async fn count_unique_ips_by_fingerprint(
217 &self,
218 fingerprint_hash: &str,
219 window_days: i64,
220 ) -> Result<i64> {
221 behavioral_queries::count_unique_ips_by_fingerprint(
222 &self.pool,
223 fingerprint_hash,
224 window_days,
225 )
226 .await
227 }
228
229 pub async fn count_engagement_events_by_fingerprint(
230 &self,
231 fingerprint_hash: &str,
232 window_days: i64,
233 ) -> Result<i64> {
234 behavioral_queries::count_engagement_events_by_fingerprint(
235 &self.pool,
236 fingerprint_hash,
237 window_days,
238 )
239 .await
240 }
241
242 pub async fn get_session_starts_by_fingerprint(
243 &self,
244 fingerprint_hash: &str,
245 window_days: i64,
246 ) -> Result<Vec<DateTime<Utc>>> {
247 behavioral_queries::get_session_starts_by_fingerprint(
248 &self.pool,
249 fingerprint_hash,
250 window_days,
251 )
252 .await
253 }
254
255 pub async fn get_session_velocity(
256 &self,
257 session_id: &SessionId,
258 ) -> Result<(Option<i64>, Option<i64>)> {
259 behavioral_queries::get_session_velocity(&self.pool, session_id).await
260 }
261}