Skip to main content

systemprompt_users/
extension.rs

1use systemprompt_extension::prelude::*;
2
3const MIGRATION_001_UTM_CONTENT_TERM: &str = r"
4ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_source VARCHAR(100);
5ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_medium VARCHAR(100);
6ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_campaign VARCHAR(100);
7ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_content VARCHAR(100);
8ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_term VARCHAR(100);
9";
10
11const MIGRATION_002_IS_AI_CRAWLER: &str = r"
12ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS is_ai_crawler BOOLEAN NOT NULL DEFAULT false;
13CREATE INDEX IF NOT EXISTS idx_user_sessions_is_ai_crawler ON user_sessions(is_ai_crawler) WHERE is_ai_crawler = true;
14UPDATE user_sessions
15SET is_ai_crawler = true,
16    is_bot = false
17WHERE is_ai_crawler = false
18  AND user_agent IS NOT NULL
19  AND (
20        user_agent ILIKE '%NotebookLM%'
21     OR user_agent ILIKE '%Gemini-Deep-Research%'
22     OR user_agent ILIKE '%Grammarly%'
23     OR user_agent ILIKE '%ChatGPT-User%'
24     OR user_agent ILIKE '%OAI-SearchBot%'
25     OR user_agent ILIKE '%GPTBot%'
26     OR user_agent ILIKE '%PerplexityBot%'
27     OR user_agent ILIKE '%Perplexity-User%'
28     OR user_agent ILIKE '%ClaudeBot%'
29     OR user_agent ILIKE '%Claude-User%'
30     OR user_agent ILIKE '%Claude-Web%'
31     OR user_agent ILIKE '%anthropic-ai%'
32     OR user_agent ILIKE '%Applebot-Extended%'
33     OR user_agent ILIKE '%CCBot%'
34     OR user_agent ILIKE '%Bytespider%'
35     OR user_agent ILIKE '%Amazonbot%'
36     OR user_agent ILIKE '%YouBot%'
37     OR user_agent ILIKE '%Diffbot%'
38     OR user_agent ILIKE '%cohere-ai%'
39  );
40";
41
42#[derive(Debug, Clone, Copy, Default)]
43pub struct UsersExtension;
44
45impl Extension for UsersExtension {
46    fn metadata(&self) -> ExtensionMetadata {
47        ExtensionMetadata {
48            id: "users",
49            name: "Users",
50            version: env!("CARGO_PKG_VERSION"),
51        }
52    }
53
54    fn migration_weight(&self) -> u32 {
55        10
56    }
57
58    fn is_required(&self) -> bool {
59        true
60    }
61
62    fn schemas(&self) -> Vec<SchemaDefinition> {
63        vec![
64            SchemaDefinition::inline("users", include_str!("../schema/users.sql"))
65                .with_required_columns(vec![
66                    "id".into(),
67                    "name".into(),
68                    "email".into(),
69                    "created_at".into(),
70                ]),
71            SchemaDefinition::inline("user_sessions", include_str!("../schema/user_sessions.sql"))
72                .with_required_columns(vec!["session_id".into(), "started_at".into()]),
73            SchemaDefinition::inline("banned_ips", include_str!("../schema/banned_ips.sql"))
74                .with_required_columns(vec![
75                    "ip_address".into(),
76                    "reason".into(),
77                    "banned_at".into(),
78                ]),
79            SchemaDefinition::inline(
80                "session_analytics_views",
81                include_str!("../schema/session_analytics_views.sql"),
82            ),
83            SchemaDefinition::inline(
84                "referrer_analytics_views",
85                include_str!("../schema/referrer_analytics_views.sql"),
86            ),
87            SchemaDefinition::inline(
88                "bot_analytics_views",
89                include_str!("../schema/bot_analytics_views.sql"),
90            ),
91            SchemaDefinition::inline("user_api_keys", include_str!("../schema/user_api_keys.sql"))
92                .with_required_columns(vec![
93                    "id".into(),
94                    "user_id".into(),
95                    "key_prefix".into(),
96                    "key_hash".into(),
97                ]),
98            SchemaDefinition::inline(
99                "user_device_certs",
100                include_str!("../schema/user_device_certs.sql"),
101            )
102            .with_required_columns(vec![
103                "id".into(),
104                "user_id".into(),
105                "fingerprint".into(),
106                "label".into(),
107            ]),
108        ]
109    }
110
111    fn dependencies(&self) -> Vec<&'static str> {
112        vec![]
113    }
114
115    fn migrations(&self) -> Vec<Migration> {
116        vec![
117            Migration::new(
118                1,
119                "add_user_sessions_utm_content_term",
120                MIGRATION_001_UTM_CONTENT_TERM,
121            ),
122            Migration::new(
123                2,
124                "add_user_sessions_is_ai_crawler",
125                MIGRATION_002_IS_AI_CRAWLER,
126            ),
127        ]
128    }
129}
130
131register_extension!(UsersExtension);