Skip to main content

systemprompt_analytics/repository/
traffic.rs

1use anyhow::Result;
2use chrono::{DateTime, Utc};
3use sqlx::PgPool;
4use std::sync::Arc;
5use systemprompt_database::DbPool;
6
7use crate::models::cli::{BotTotalsRow, BotTypeRow, DeviceRow, GeoRow, TrafficSourceRow};
8
9#[derive(Debug)]
10pub struct TrafficAnalyticsRepository {
11    pool: Arc<PgPool>,
12}
13
14impl TrafficAnalyticsRepository {
15    pub fn new(db: &DbPool) -> Result<Self> {
16        let pool = db.pool_arc()?;
17        Ok(Self { pool })
18    }
19
20    pub async fn get_sources(
21        &self,
22        start: DateTime<Utc>,
23        end: DateTime<Utc>,
24        limit: i64,
25    ) -> Result<Vec<TrafficSourceRow>> {
26        sqlx::query_as!(
27            TrafficSourceRow,
28            r#"
29            SELECT
30                COALESCE(referrer_source, 'direct') as "source",
31                COUNT(*)::bigint as "count!"
32            FROM user_sessions
33            WHERE started_at >= $1 AND started_at < $2
34              AND is_bot = false AND is_behavioral_bot = false AND is_scanner = false
35            GROUP BY referrer_source
36            ORDER BY COUNT(*) DESC
37            LIMIT $3
38            "#,
39            start,
40            end,
41            limit
42        )
43        .fetch_all(&*self.pool)
44        .await
45        .map_err(Into::into)
46    }
47
48    pub async fn get_geo_breakdown(
49        &self,
50        start: DateTime<Utc>,
51        end: DateTime<Utc>,
52        limit: i64,
53    ) -> Result<Vec<GeoRow>> {
54        sqlx::query_as!(
55            GeoRow,
56            r#"
57            SELECT
58                COALESCE(country, 'Unknown') as "country",
59                COUNT(*)::bigint as "count!"
60            FROM user_sessions
61            WHERE started_at >= $1 AND started_at < $2
62              AND is_bot = false AND is_behavioral_bot = false AND is_scanner = false
63            GROUP BY country
64            ORDER BY COUNT(*) DESC
65            LIMIT $3
66            "#,
67            start,
68            end,
69            limit
70        )
71        .fetch_all(&*self.pool)
72        .await
73        .map_err(Into::into)
74    }
75
76    pub async fn get_device_breakdown(
77        &self,
78        start: DateTime<Utc>,
79        end: DateTime<Utc>,
80        limit: i64,
81    ) -> Result<Vec<DeviceRow>> {
82        sqlx::query_as!(
83            DeviceRow,
84            r#"
85            SELECT
86                COALESCE(device_type, 'unknown') as "device",
87                COALESCE(browser, 'unknown') as "browser",
88                COUNT(*)::bigint as "count!"
89            FROM user_sessions
90            WHERE started_at >= $1 AND started_at < $2
91              AND is_bot = false AND is_behavioral_bot = false AND is_scanner = false
92            GROUP BY device_type, browser
93            ORDER BY COUNT(*) DESC
94            LIMIT $3
95            "#,
96            start,
97            end,
98            limit
99        )
100        .fetch_all(&*self.pool)
101        .await
102        .map_err(Into::into)
103    }
104
105    pub async fn get_bot_totals(
106        &self,
107        start: DateTime<Utc>,
108        end: DateTime<Utc>,
109    ) -> Result<BotTotalsRow> {
110        sqlx::query_as!(
111            BotTotalsRow,
112            r#"
113            SELECT
114                COUNT(*) FILTER (WHERE is_bot = false AND is_behavioral_bot = false AND is_scanner = false)::bigint as "human!",
115                COUNT(*) FILTER (WHERE is_bot = true OR is_behavioral_bot = true OR is_scanner = true)::bigint as "bot!"
116            FROM user_sessions
117            WHERE started_at >= $1 AND started_at < $2
118            "#,
119            start,
120            end
121        )
122        .fetch_one(&*self.pool)
123        .await
124        .map_err(Into::into)
125    }
126
127    pub async fn get_bot_breakdown(
128        &self,
129        start: DateTime<Utc>,
130        end: DateTime<Utc>,
131    ) -> Result<Vec<BotTypeRow>> {
132        sqlx::query_as!(
133            BotTypeRow,
134            r#"
135            SELECT
136                COALESCE(
137                    CASE
138                        WHEN user_agent ILIKE '%googlebot%' THEN 'Googlebot'
139                        WHEN user_agent ILIKE '%bingbot%' THEN 'Bingbot'
140                        WHEN user_agent ILIKE '%chatgpt%' THEN 'ChatGPT'
141                        WHEN user_agent ILIKE '%claude%' THEN 'Claude'
142                        WHEN user_agent ILIKE '%perplexity%' THEN 'Perplexity'
143                        ELSE 'Other'
144                    END,
145                    'Unknown'
146                ) as "bot_type",
147                COUNT(*)::bigint as "count!"
148            FROM user_sessions
149            WHERE started_at >= $1 AND started_at < $2
150              AND (is_bot = true OR is_behavioral_bot = true OR is_scanner = true)
151            GROUP BY 1
152            ORDER BY COUNT(*) DESC
153            "#,
154            start,
155            end
156        )
157        .fetch_all(&*self.pool)
158        .await
159        .map_err(Into::into)
160    }
161}