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_as_behavioral_bot(&self, session_id: &SessionId, reason: &str) -> Result<()> {
78        behavioral::mark_as_behavioral_bot(&self.write_pool, session_id, reason).await
79    }
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
94    pub async fn cleanup_inactive(&self, inactive_hours: i32) -> Result<u64> {
95        mutations::cleanup_inactive(&self.write_pool, inactive_hours).await
96    }
97
98    pub async fn migrate_user_sessions(
99        &self,
100        old_user_id: &UserId,
101        new_user_id: &UserId,
102    ) -> Result<u64> {
103        mutations::migrate_user_sessions(&self.write_pool, old_user_id, new_user_id).await
104    }
105
106    pub async fn create_session(&self, params: &CreateSessionParams<'_>) -> Result<()> {
107        mutations::create_session(&self.write_pool, params).await
108    }
109
110    pub async fn find_recent_by_fingerprint(
111        &self,
112        fingerprint_hash: &str,
113        max_age_seconds: i64,
114    ) -> Result<Option<SessionRecord>> {
115        queries::find_recent_by_fingerprint(&self.pool, fingerprint_hash, max_age_seconds).await
116    }
117
118    pub async fn exists(&self, session_id: &SessionId) -> Result<bool> {
119        queries::exists(&self.pool, session_id).await
120    }
121
122    pub async fn increment_ai_usage(
123        &self,
124        session_id: &SessionId,
125        tokens: i32,
126        cost_microdollars: i64,
127    ) -> Result<()> {
128        mutations::increment_ai_usage(&self.write_pool, session_id, tokens, cost_microdollars).await
129    }
130
131    pub async fn update_behavioral_detection(
132        &self,
133        session_id: &SessionId,
134        score: i32,
135        is_behavioral_bot: bool,
136        reason: Option<&str>,
137    ) -> Result<()> {
138        behavioral::update_behavioral_detection(
139            &self.write_pool,
140            session_id,
141            score,
142            is_behavioral_bot,
143            reason,
144        )
145        .await
146    }
147
148    pub async fn escalate_throttle(&self, session_id: &SessionId, new_level: i32) -> Result<()> {
149        mutations::escalate_throttle(&self.write_pool, session_id, new_level).await
150    }
151
152    pub async fn get_throttle_level(&self, session_id: &SessionId) -> Result<i32> {
153        queries::get_throttle_level(&self.pool, session_id).await
154    }
155
156    pub async fn count_sessions_by_fingerprint(
157        &self,
158        fingerprint_hash: &str,
159        window_hours: i64,
160    ) -> Result<i64> {
161        queries::count_sessions_by_fingerprint(&self.pool, fingerprint_hash, window_hours).await
162    }
163
164    pub async fn get_endpoint_sequence(&self, session_id: &SessionId) -> Result<Vec<String>> {
165        queries::get_endpoint_sequence(&self.pool, session_id).await
166    }
167
168    pub async fn get_request_timestamps(
169        &self,
170        session_id: &SessionId,
171    ) -> Result<Vec<DateTime<Utc>>> {
172        queries::get_request_timestamps(&self.pool, session_id).await
173    }
174
175    pub async fn get_total_content_pages(&self) -> Result<i64> {
176        queries::get_total_content_pages(&self.pool).await
177    }
178
179    pub async fn get_session_for_behavioral_analysis(
180        &self,
181        session_id: &SessionId,
182    ) -> Result<Option<SessionBehavioralData>> {
183        queries::get_session_for_behavioral_analysis(&self.pool, session_id).await
184    }
185
186    pub async fn has_analytics_events(&self, session_id: &SessionId) -> Result<bool> {
187        queries::has_analytics_events(&self.pool, session_id).await
188    }
189
190    pub async fn get_session_velocity(
191        &self,
192        session_id: &SessionId,
193    ) -> Result<(Option<i64>, Option<i64>)> {
194        queries::get_session_velocity(&self.pool, session_id).await
195    }
196}
systemprompt_analytics/repository/session/mod.rs

systemprompt_analytics/repository/session/
mod.rs
