seed.rs

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