systemprompt-users 0.6.0

use systemprompt_extension::prelude::*;

const MIGRATION_001_UTM_CONTENT_TERM: &str = r"
ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_source VARCHAR(100);
ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_medium VARCHAR(100);
ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_campaign VARCHAR(100);
ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_content VARCHAR(100);
ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS utm_term VARCHAR(100);
";

const MIGRATION_002_IS_AI_CRAWLER: &str = r"
ALTER TABLE user_sessions ADD COLUMN IF NOT EXISTS is_ai_crawler BOOLEAN NOT NULL DEFAULT false;
CREATE INDEX IF NOT EXISTS idx_user_sessions_is_ai_crawler ON user_sessions(is_ai_crawler) WHERE is_ai_crawler = true;
UPDATE user_sessions
SET is_ai_crawler = true,
    is_bot = false
WHERE is_ai_crawler = false
  AND user_agent IS NOT NULL
  AND (
        user_agent ILIKE '%NotebookLM%'
     OR user_agent ILIKE '%Gemini-Deep-Research%'
     OR user_agent ILIKE '%Grammarly%'
     OR user_agent ILIKE '%ChatGPT-User%'
     OR user_agent ILIKE '%OAI-SearchBot%'
     OR user_agent ILIKE '%GPTBot%'
     OR user_agent ILIKE '%PerplexityBot%'
     OR user_agent ILIKE '%Perplexity-User%'
     OR user_agent ILIKE '%ClaudeBot%'
     OR user_agent ILIKE '%Claude-User%'
     OR user_agent ILIKE '%Claude-Web%'
     OR user_agent ILIKE '%anthropic-ai%'
     OR user_agent ILIKE '%Applebot-Extended%'
     OR user_agent ILIKE '%CCBot%'
     OR user_agent ILIKE '%Bytespider%'
     OR user_agent ILIKE '%Amazonbot%'
     OR user_agent ILIKE '%YouBot%'
     OR user_agent ILIKE '%Diffbot%'
     OR user_agent ILIKE '%cohere-ai%'
  );
";

#[derive(Debug, Clone, Copy, Default)]
pub struct UsersExtension;

impl Extension for UsersExtension {
    fn metadata(&self) -> ExtensionMetadata {
        ExtensionMetadata {
            id: "users",
            name: "Users",
            version: env!("CARGO_PKG_VERSION"),
        }
    }

    fn migration_weight(&self) -> u32 {
        100
    }

    fn is_required(&self) -> bool {
        true
    }

    fn schemas(&self) -> Vec<SchemaDefinition> {
        vec![
            SchemaDefinition::inline("users", include_str!("../schema/users.sql"))
                .with_required_columns(vec![
                    "id".into(),
                    "name".into(),
                    "email".into(),
                    "created_at".into(),
                ]),
            SchemaDefinition::inline("user_sessions", include_str!("../schema/user_sessions.sql"))
                .with_required_columns(vec!["session_id".into(), "started_at".into()]),
            SchemaDefinition::inline("banned_ips", include_str!("../schema/banned_ips.sql"))
                .with_required_columns(vec![
                    "ip_address".into(),
                    "reason".into(),
                    "banned_at".into(),
                ]),
            SchemaDefinition::inline(
                "session_analytics_views",
                include_str!("../schema/session_analytics_views.sql"),
            ),
            SchemaDefinition::inline(
                "referrer_analytics_views",
                include_str!("../schema/referrer_analytics_views.sql"),
            ),
            SchemaDefinition::inline(
                "bot_analytics_views",
                include_str!("../schema/bot_analytics_views.sql"),
            ),
            SchemaDefinition::inline("user_api_keys", include_str!("../schema/user_api_keys.sql"))
                .with_required_columns(vec![
                    "id".into(),
                    "user_id".into(),
                    "key_prefix".into(),
                    "key_hash".into(),
                ]),
            SchemaDefinition::inline(
                "user_device_certs",
                include_str!("../schema/user_device_certs.sql"),
            )
            .with_required_columns(vec![
                "id".into(),
                "user_id".into(),
                "fingerprint".into(),
                "label".into(),
            ]),
        ]
    }

    fn dependencies(&self) -> Vec<&'static str> {
        vec![]
    }

    fn migrations(&self) -> Vec<Migration> {
        vec![
            Migration::new(
                1,
                "add_user_sessions_utm_content_term",
                MIGRATION_001_UTM_CONTENT_TERM,
            ),
            Migration::new(
                2,
                "add_user_sessions_is_ai_crawler",
                MIGRATION_002_IS_AI_CRAWLER,
            ),
        ]
    }
}

register_extension!(UsersExtension);