seed.rs

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