systemprompt_users/repository/session/
mod.rs1mod behavioral;
7mod behavioral_queries;
8mod fingerprint;
9mod geo;
10mod mutations;
11mod queries;
12mod snapshot_row;
13
14use crate::Result;
15use chrono::{DateTime, Utc};
16use sqlx::PgPool;
17use std::sync::Arc;
18use systemprompt_database::DbPool;
19use systemprompt_identifiers::{SessionId, UserId};
20use systemprompt_models::analytics::BehavioralBotReason;
21use systemprompt_traits::session_store::{
22 ActiveSessionLookup, CreateSessionParams, SessionBehavioralData, SessionRecord,
23 SessionSnapshot as AnalyticsSession,
24};
25#[derive(Clone, Debug)]
26pub struct SessionRepository {
27 pool: Arc<PgPool>,
28 write_pool: Arc<PgPool>,
29}
30impl SessionRepository {
31 pub fn new(db: &DbPool) -> Self {
32 Self {
33 pool: db.pool(),
34 write_pool: db.write_pool(),
35 }
36 }
37 pub async fn find_by_id(&self, session_id: &SessionId) -> Result<Option<AnalyticsSession>> {
38 queries::find_by_id(&self.write_pool, session_id).await
39 }
40 pub async fn find_active_by_id(
41 &self,
42 session_id: &SessionId,
43 ) -> Result<Option<ActiveSessionLookup>> {
44 queries::find_active_by_id(&self.write_pool, session_id).await
45 }
46 pub async fn revoke_session(&self, session_id: &SessionId) -> Result<()> {
47 mutations::revoke_session(&self.write_pool, session_id).await
48 }
49 pub async fn revoke_all_for_user(&self, user_id: &UserId) -> Result<u64> {
50 mutations::revoke_all_for_user(&self.write_pool, user_id).await
51 }
52 pub async fn find_by_fingerprint(
53 &self,
54 fingerprint_hash: &str,
55 user_id: &UserId,
56 ) -> Result<Option<AnalyticsSession>> {
57 queries::find_by_fingerprint(&self.pool, fingerprint_hash, user_id).await
58 }
59 pub async fn list_active_by_user(&self, user_id: &UserId) -> Result<Vec<AnalyticsSession>> {
60 queries::list_active_by_user(&self.pool, user_id).await
61 }
62 pub async fn increment_request_count(&self, session_id: &SessionId) -> Result<()> {
63 mutations::increment_request_count(&self.write_pool, session_id).await
64 }
65 pub async fn increment_task_count(&self, session_id: &SessionId) -> Result<()> {
66 mutations::increment_task_count(&self.write_pool, session_id).await
67 }
68 pub async fn increment_message_count(&self, session_id: &SessionId) -> Result<()> {
69 mutations::increment_message_count(&self.write_pool, session_id).await
70 }
71 pub async fn end_session(&self, session_id: &SessionId) -> Result<()> {
72 mutations::end_session(&self.write_pool, session_id).await
73 }
74 pub async fn mark_as_scanner(&self, session_id: &SessionId) -> Result<()> {
75 mutations::mark_as_scanner(&self.write_pool, session_id).await
76 }
77 pub async fn mark_converted(&self, session_id: &SessionId) -> Result<()> {
78 mutations::mark_converted(&self.write_pool, session_id).await
79 }
80 pub async fn mark_as_behavioral_bot(
81 &self,
82 session_id: &SessionId,
83 reason: BehavioralBotReason,
84 ) -> Result<()> {
85 behavioral::mark_as_behavioral_bot(&self.write_pool, session_id, reason).await
86 }
87 pub async fn mark_ghost_sessions_as_bots(&self) -> Result<u64> {
88 behavioral::mark_ghost_sessions_as_bots(&self.write_pool).await
89 }
90 pub async fn mark_no_js_sessions_as_bots(&self, engaged: &[SessionId]) -> Result<u64> {
91 let engaged: Vec<String> = engaged.iter().map(ToString::to_string).collect();
92 behavioral::mark_no_js_sessions_as_bots(&self.write_pool, &engaged).await
93 }
94 pub async fn check_and_mark_behavioral_bot(
95 &self,
96 session_id: &SessionId,
97 request_count_threshold: i32,
98 ) -> Result<bool> {
99 behavioral::check_and_mark_behavioral_bot(
100 &self.write_pool,
101 session_id,
102 request_count_threshold,
103 )
104 .await
105 }
106 pub async fn cleanup_inactive(&self, inactive_hours: i32) -> Result<u64> {
107 mutations::cleanup_inactive(&self.write_pool, inactive_hours).await
108 }
109 pub async fn count_inactive(&self, inactive_hours: i32) -> Result<i64> {
110 queries::count_inactive(&self.pool, inactive_hours).await
111 }
112 pub async fn count_sessions_missing_geo(&self) -> Result<i64> {
113 mutations::count_sessions_missing_geo(&self.pool).await
114 }
115 pub async fn migrate_user_sessions(
116 &self,
117 old_user_id: &UserId,
118 new_user_id: &UserId,
119 ) -> Result<u64> {
120 mutations::migrate_user_sessions(&self.write_pool, old_user_id, new_user_id).await
121 }
122 pub async fn create_session(&self, params: &CreateSessionParams<'_>) -> Result<()> {
123 mutations::create_session(&self.write_pool, params).await
124 }
125 pub async fn find_recent_anonymous_by_fingerprint(
126 &self,
127 fingerprint_hash: &str,
128 max_age_seconds: i64,
129 ) -> Result<Option<SessionRecord>> {
130 queries::find_recent_anonymous_by_fingerprint(
131 &self.write_pool,
132 fingerprint_hash,
133 max_age_seconds,
134 )
135 .await
136 }
137 pub async fn increment_ai_usage(
138 &self,
139 session_id: &SessionId,
140 tokens: i32,
141 cost_microdollars: i64,
142 ) -> Result<()> {
143 mutations::increment_ai_usage(&self.write_pool, session_id, tokens, cost_microdollars).await
144 }
145 pub async fn update_behavioral_detection(
146 &self,
147 session_id: &SessionId,
148 score: i32,
149 is_behavioral_bot: bool,
150 reason: Option<&str>,
151 ) -> Result<()> {
152 behavioral::update_behavioral_detection(
153 &self.write_pool,
154 session_id,
155 score,
156 is_behavioral_bot,
157 reason,
158 )
159 .await
160 }
161 pub async fn count_sessions_by_fingerprint(
162 &self,
163 fingerprint_hash: &str,
164 window_hours: i64,
165 ) -> Result<i64> {
166 behavioral_queries::count_sessions_by_fingerprint(
167 &self.write_pool,
168 fingerprint_hash,
169 window_hours,
170 )
171 .await
172 }
173 pub async fn get_session_for_behavioral_analysis(
174 &self,
175 session_id: &SessionId,
176 ) -> Result<Option<SessionBehavioralData>> {
177 behavioral_queries::get_session_for_behavioral_analysis(&self.write_pool, session_id).await
178 }
179 pub async fn count_unique_ips_by_fingerprint(
180 &self,
181 fingerprint_hash: &str,
182 window_days: i64,
183 ) -> Result<i64> {
184 behavioral_queries::count_unique_ips_by_fingerprint(
185 &self.write_pool,
186 fingerprint_hash,
187 window_days,
188 )
189 .await
190 }
191 pub async fn get_session_starts_by_fingerprint(
192 &self,
193 fingerprint_hash: &str,
194 window_days: i64,
195 ) -> Result<Vec<DateTime<Utc>>> {
196 behavioral_queries::get_session_starts_by_fingerprint(
197 &self.write_pool,
198 fingerprint_hash,
199 window_days,
200 )
201 .await
202 }
203 pub async fn get_session_velocity(
204 &self,
205 session_id: &SessionId,
206 ) -> Result<(Option<i64>, Option<i64>)> {
207 behavioral_queries::get_session_velocity(&self.write_pool, session_id).await
208 }
209}