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