systemprompt_users/repository/
session_traffic.rs1use std::sync::Arc;
8
9use async_trait::async_trait;
10use chrono::{DateTime, Utc};
11use sqlx::PgPool;
12use systemprompt_database::DbPool;
13use systemprompt_models::analytics::TrafficClass;
14use systemprompt_traits::{RepositoryError, TrafficClassCounts, TrafficClassTotals};
15
16#[derive(Debug, Clone)]
17pub struct SessionTrafficRollupRepository {
18 pool: Arc<PgPool>,
19}
20
21impl SessionTrafficRollupRepository {
22 pub fn new(db: &DbPool) -> Self {
23 Self { pool: db.pool() }
24 }
25
26 pub async fn rebuild(&self) -> Result<i64, RepositoryError> {
27 let written =
28 sqlx::query_scalar!(r#"SELECT session_traffic_rollup_rebuild() AS "written!""#)
29 .fetch_one(&*self.pool)
30 .await?;
31 Ok(written)
32 }
33}
34
35#[async_trait]
36impl TrafficClassCounts for SessionTrafficRollupRepository {
37 async fn session_totals(
38 &self,
39 start: DateTime<Utc>,
40 end: DateTime<Utc>,
41 ) -> Result<TrafficClassTotals, RepositoryError> {
42 let rows = sqlx::query!(
43 r#"
44 SELECT traffic_class, COALESCE(SUM(sessions), 0)::bigint AS "sessions!"
45 FROM session_traffic_rollup_hourly
46 WHERE bucket >= date_trunc('hour', $1::timestamptz) AND bucket < $2
47 GROUP BY traffic_class
48 "#,
49 start,
50 end
51 )
52 .fetch_all(&*self.pool)
53 .await?;
54
55 let mut totals = TrafficClassTotals::default();
56 for row in rows {
57 let class = TrafficClass::parse(&row.traffic_class).ok_or_else(|| {
58 RepositoryError::decode(
59 "traffic_class",
60 std::io::Error::new(std::io::ErrorKind::InvalidData, row.traffic_class.clone()),
61 )
62 })?;
63 match class {
64 TrafficClass::Human => totals.human = row.sessions,
65 TrafficClass::Ghost => totals.ghost = row.sessions,
66 TrafficClass::Bot => totals.bot = row.sessions,
67 }
68 }
69 Ok(totals)
70 }
71}