Skip to main content

systemprompt_users/repository/
session_traffic.rs

1//! `session_traffic_rollup_hourly` reads: the write-time counter behind the
2//! traffic totals analytics reports.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use 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}