seed.rs

 1use collab::{
 2    db::{self, NewUserParams},
 3    env::load_dotenv,
 4    executor::Executor,
 5};
 6use db::{ConnectOptions, Database};
 7use serde::{de::DeserializeOwned, Deserialize};
 8use std::{fmt::Write, fs};
 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    load_dotenv().expect("failed to load .env.toml file");
20
21    let mut admin_logins = load_admins("crates/collab/.admins.default.json")
22        .expect("failed to load default admins file");
23    if let Ok(other_admins) = load_admins("./.admins.json") {
24        admin_logins.extend(other_admins);
25    }
26
27    let database_url = std::env::var("DATABASE_URL").expect("missing DATABASE_URL env var");
28    let db = Database::new(ConnectOptions::new(database_url), Executor::Production)
29        .await
30        .expect("failed to connect to postgres database");
31    let client = reqwest::Client::new();
32
33    // Create admin users for all of the users in `.admins.toml` or `.admins.default.toml`.
34    for admin_login in admin_logins {
35        let user = fetch_github::<GitHubUser>(
36            &client,
37            &format!("https://api.github.com/users/{admin_login}"),
38        )
39        .await;
40        db.create_user(
41            &user.email.unwrap_or(format!("{admin_login}@example.com")),
42            true,
43            NewUserParams {
44                github_login: user.login,
45                github_user_id: user.id,
46            },
47        )
48        .await
49        .expect("failed to create admin user");
50    }
51
52    // Fetch 100 other random users from GitHub and insert them into the database.
53    let mut user_count = db
54        .get_all_users(0, 200)
55        .await
56        .expect("failed to load users from db")
57        .len();
58    let mut last_user_id = None;
59    while user_count < 100 {
60        let mut uri = "https://api.github.com/users?per_page=100".to_string();
61        if let Some(last_user_id) = last_user_id {
62            write!(&mut uri, "&since={}", last_user_id).unwrap();
63        }
64        let users = fetch_github::<Vec<GitHubUser>>(&client, &uri).await;
65
66        for github_user in users {
67            last_user_id = Some(github_user.id);
68            user_count += 1;
69            db.get_or_create_user_by_github_account(
70                &github_user.login,
71                Some(github_user.id),
72                github_user.email.as_deref(),
73                None,
74            )
75            .await
76            .expect("failed to insert user");
77        }
78    }
79}
80
81fn load_admins(path: &str) -> anyhow::Result<Vec<String>> {
82    let file_content = fs::read_to_string(path)?;
83    Ok(serde_json::from_str(&file_content)?)
84}
85
86async fn fetch_github<T: DeserializeOwned>(client: &reqwest::Client, url: &str) -> T {
87    let response = client
88        .get(url)
89        .header("user-agent", "zed")
90        .send()
91        .await
92        .unwrap_or_else(|_| panic!("failed to fetch '{}'", url));
93    response
94        .json()
95        .await
96        .unwrap_or_else(|_| panic!("failed to deserialize github user from '{}'", url))
97}