systemprompt_analytics/repository/session/
mod.rs

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

systemprompt_analytics/repository/session/
mod.rs
