seed.rs

  1use collab2::{db, executor::Executor};
  2use db::{ConnectOptions, Database};
  3use serde::{de::DeserializeOwned, Deserialize};
  4use std::fmt::Write;
  5
  6#[derive(Debug, Deserialize)]
  7struct GitHubUser {
  8    id: i32,
  9    login: String,
 10    email: Option<String>,
 11}
 12
 13#[tokio::main]
 14async fn main() {
 15    let database_url = std::env::var("DATABASE_URL").expect("missing DATABASE_URL env var");
 16    let db = Database::new(ConnectOptions::new(database_url), Executor::Production)
 17        .await
 18        .expect("failed to connect to postgres database");
 19    let github_token = std::env::var("GITHUB_TOKEN").expect("missing GITHUB_TOKEN env var");
 20    let client = reqwest::Client::new();
 21
 22    let mut current_user =
 23        fetch_github::<GitHubUser>(&client, &github_token, "https://api.github.com/user").await;
 24    current_user
 25        .email
 26        .get_or_insert_with(|| "placeholder@example.com".to_string());
 27    let staff_users = fetch_github::<Vec<GitHubUser>>(
 28        &client,
 29        &github_token,
 30        "https://api.github.com/orgs/zed-industries/teams/staff/members",
 31    )
 32    .await;
 33
 34    let mut zed_users = Vec::new();
 35    zed_users.push((current_user, true));
 36    zed_users.extend(staff_users.into_iter().map(|user| (user, true)));
 37
 38    let user_count = db
 39        .get_all_users(0, 200)
 40        .await
 41        .expect("failed to load users from db")
 42        .len();
 43    if user_count < 100 {
 44        let mut last_user_id = None;
 45        for _ in 0..10 {
 46            let mut uri = "https://api.github.com/users?per_page=100".to_string();
 47            if let Some(last_user_id) = last_user_id {
 48                write!(&mut uri, "&since={}", last_user_id).unwrap();
 49            }
 50            let users = fetch_github::<Vec<GitHubUser>>(&client, &github_token, &uri).await;
 51            if let Some(last_user) = users.last() {
 52                last_user_id = Some(last_user.id);
 53                zed_users.extend(users.into_iter().map(|user| (user, false)));
 54            } else {
 55                break;
 56            }
 57        }
 58    }
 59
 60    for (github_user, admin) in zed_users {
 61        if db
 62            .get_user_by_github_login(&github_user.login)
 63            .await
 64            .expect("failed to fetch user")
 65            .is_none()
 66        {
 67            if admin {
 68                db.create_user(
 69                    &format!("{}@zed.dev", github_user.login),
 70                    admin,
 71                    db::NewUserParams {
 72                        github_login: github_user.login,
 73                        github_user_id: github_user.id,
 74                    },
 75                )
 76                .await
 77                .expect("failed to insert user");
 78            } else {
 79                db.get_or_create_user_by_github_account(
 80                    &github_user.login,
 81                    Some(github_user.id),
 82                    github_user.email.as_deref(),
 83                )
 84                .await
 85                .expect("failed to insert user");
 86            }
 87        }
 88    }
 89}
 90
 91async fn fetch_github<T: DeserializeOwned>(
 92    client: &reqwest::Client,
 93    access_token: &str,
 94    url: &str,
 95) -> T {
 96    let response = client
 97        .get(url)
 98        .bearer_auth(&access_token)
 99        .header("user-agent", "zed")
100        .send()
101        .await
102        .expect(&format!("failed to fetch '{}'", url));
103    response
104        .json()
105        .await
106        .expect(&format!("failed to deserialize github user from '{}'", url))
107}