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