seed.rs

  1use crate::db::{self, ChannelRole, NewUserParams};
  2
  3use anyhow::Context;
  4use chrono::{DateTime, Utc};
  5use db::Database;
  6use serde::{de::DeserializeOwned, Deserialize};
  7use std::{fmt::Write, fs, path::Path};
  8
  9use crate::Config;
 10
 11#[derive(Debug, Deserialize)]
 12struct GitHubUser {
 13    id: i32,
 14    login: String,
 15    email: Option<String>,
 16    created_at: DateTime<Utc>,
 17}
 18
 19#[derive(Deserialize)]
 20struct SeedConfig {
 21    // Which users to create as admins.
 22    admins: Vec<String>,
 23    // Which channels to create (all admins are invited to all channels)
 24    channels: Vec<String>,
 25    // Number of random users to create from the Github API
 26    number_of_users: Option<usize>,
 27}
 28
 29pub async fn seed(config: &Config, db: &Database, force: bool) -> anyhow::Result<()> {
 30    let client = reqwest::Client::new();
 31
 32    if !db.get_all_users(0, 1).await?.is_empty() && !force {
 33        return Ok(());
 34    }
 35
 36    let seed_path = config
 37        .seed_path
 38        .as_ref()
 39        .context("called seed with no SEED_PATH")?;
 40
 41    let seed_config = load_admins(seed_path)
 42        .context(format!("failed to load {}", seed_path.to_string_lossy()))?;
 43
 44    let mut first_user = None;
 45    let mut others = vec![];
 46
 47    for admin_login in seed_config.admins {
 48        let user = fetch_github::<GitHubUser>(
 49            &client,
 50            &format!("https://api.github.com/users/{admin_login}"),
 51        )
 52        .await;
 53        let user = db
 54            .create_user(
 55                &user.email.unwrap_or(format!("{admin_login}@example.com")),
 56                true,
 57                NewUserParams {
 58                    github_login: user.login,
 59                    github_user_id: user.id,
 60                },
 61            )
 62            .await
 63            .context("failed to create admin user")?;
 64        if first_user.is_none() {
 65            first_user = Some(user.user_id);
 66        } else {
 67            others.push(user.user_id)
 68        }
 69    }
 70
 71    for channel in seed_config.channels {
 72        let (channel, _) = db
 73            .create_channel(&channel, None, first_user.unwrap())
 74            .await
 75            .context("failed to create channel")?;
 76
 77        for user_id in &others {
 78            db.invite_channel_member(
 79                channel.id,
 80                *user_id,
 81                first_user.unwrap(),
 82                ChannelRole::Admin,
 83            )
 84            .await
 85            .context("failed to add user to channel")?;
 86        }
 87    }
 88
 89    if let Some(number_of_users) = seed_config.number_of_users {
 90        // Fetch 100 other random users from GitHub and insert them into the database
 91        // (for testing autocompleters, etc.)
 92        let mut user_count = db
 93            .get_all_users(0, 200)
 94            .await
 95            .expect("failed to load users from db")
 96            .len();
 97        let mut last_user_id = None;
 98        while user_count < number_of_users {
 99            let mut uri = "https://api.github.com/users?per_page=100".to_string();
100            if let Some(last_user_id) = last_user_id {
101                write!(&mut uri, "&since={}", last_user_id).unwrap();
102            }
103            let users = fetch_github::<Vec<GitHubUser>>(&client, &uri).await;
104
105            for github_user in users {
106                last_user_id = Some(github_user.id);
107                user_count += 1;
108                db.get_or_create_user_by_github_account(
109                    &github_user.login,
110                    Some(github_user.id),
111                    github_user.email.as_deref(),
112                    Some(github_user.created_at),
113                    None,
114                )
115                .await
116                .expect("failed to insert user");
117            }
118        }
119    }
120
121    Ok(())
122}
123
124fn load_admins(path: impl AsRef<Path>) -> anyhow::Result<SeedConfig> {
125    let file_content = fs::read_to_string(path)?;
126    Ok(serde_json::from_str(&file_content)?)
127}
128
129async fn fetch_github<T: DeserializeOwned>(client: &reqwest::Client, url: &str) -> T {
130    let response = client
131        .get(url)
132        .header("user-agent", "zed")
133        .send()
134        .await
135        .unwrap_or_else(|_| panic!("failed to fetch '{}'", url));
136    response
137        .json()
138        .await
139        .unwrap_or_else(|_| panic!("failed to deserialize github user from '{}'", url))
140}