seed.rs

  1use collab::{Error, Result};
  2use db::{Db, PostgresDb, UserId};
  3use rand::prelude::*;
  4use serde::{de::DeserializeOwned, Deserialize};
  5use std::fmt::Write;
  6use time::{Duration, OffsetDateTime};
  7
  8#[allow(unused)]
  9#[path = "../db.rs"]
 10mod db;
 11
 12#[derive(Debug, Deserialize)]
 13struct GitHubUser {
 14    id: i32,
 15    login: String,
 16    email: Option<String>,
 17}
 18
 19#[tokio::main]
 20async fn main() {
 21    let mut rng = StdRng::from_entropy();
 22    let database_url = std::env::var("DATABASE_URL").expect("missing DATABASE_URL env var");
 23    let db = PostgresDb::new(&database_url, 5)
 24        .await
 25        .expect("failed to connect to postgres database");
 26    let github_token = std::env::var("GITHUB_TOKEN").expect("missing GITHUB_TOKEN env var");
 27    let client = reqwest::Client::new();
 28
 29    let mut current_user =
 30        fetch_github::<GitHubUser>(&client, &github_token, "https://api.github.com/user").await;
 31    current_user
 32        .email
 33        .get_or_insert_with(|| "placeholder@example.com".to_string());
 34    let staff_users = fetch_github::<Vec<GitHubUser>>(
 35        &client,
 36        &github_token,
 37        "https://api.github.com/orgs/zed-industries/teams/staff/members",
 38    )
 39    .await;
 40
 41    let mut zed_users = Vec::new();
 42    zed_users.push((current_user, true));
 43    zed_users.extend(staff_users.into_iter().map(|user| (user, true)));
 44
 45    let user_count = db
 46        .get_all_users(0, 200)
 47        .await
 48        .expect("failed to load users from db")
 49        .len();
 50    if user_count < 100 {
 51        let mut last_user_id = None;
 52        for _ in 0..10 {
 53            let mut uri = "https://api.github.com/users?per_page=100".to_string();
 54            if let Some(last_user_id) = last_user_id {
 55                write!(&mut uri, "&since={}", last_user_id).unwrap();
 56            }
 57            let users = fetch_github::<Vec<GitHubUser>>(&client, &github_token, &uri).await;
 58            if let Some(last_user) = users.last() {
 59                last_user_id = Some(last_user.id);
 60                zed_users.extend(users.into_iter().map(|user| (user, false)));
 61            } else {
 62                break;
 63            }
 64        }
 65    }
 66
 67    let mut zed_user_ids = Vec::<UserId>::new();
 68    for (github_user, admin) in zed_users {
 69        if let Some(user) = db
 70            .get_user_by_github_account(&github_user.login, Some(github_user.id))
 71            .await
 72            .expect("failed to fetch user")
 73        {
 74            zed_user_ids.push(user.id);
 75        } else if let Some(email) = &github_user.email {
 76            zed_user_ids.push(
 77                db.create_user(
 78                    email,
 79                    admin,
 80                    db::NewUserParams {
 81                        github_login: github_user.login,
 82                        github_user_id: github_user.id,
 83                        invite_count: 5,
 84                    },
 85                )
 86                .await
 87                .expect("failed to insert user")
 88                .user_id,
 89            );
 90        } else if admin {
 91            zed_user_ids.push(
 92                db.create_user(
 93                    &format!("{}@zed.dev", github_user.login),
 94                    admin,
 95                    db::NewUserParams {
 96                        github_login: github_user.login,
 97                        github_user_id: github_user.id,
 98                        invite_count: 5,
 99                    },
100                )
101                .await
102                .expect("failed to insert user")
103                .user_id,
104            );
105        }
106    }
107
108    let zed_org_id = if let Some(org) = db
109        .find_org_by_slug("zed")
110        .await
111        .expect("failed to fetch org")
112    {
113        org.id
114    } else {
115        db.create_org("Zed", "zed")
116            .await
117            .expect("failed to insert org")
118    };
119
120    let general_channel_id = if let Some(channel) = db
121        .get_org_channels(zed_org_id)
122        .await
123        .expect("failed to fetch channels")
124        .iter()
125        .find(|c| c.name == "General")
126    {
127        channel.id
128    } else {
129        let channel_id = db
130            .create_org_channel(zed_org_id, "General")
131            .await
132            .expect("failed to insert channel");
133
134        let now = OffsetDateTime::now_utc();
135        let max_seconds = Duration::days(100).as_seconds_f64();
136        let mut timestamps = (0..1000)
137            .map(|_| now - Duration::seconds_f64(rng.gen_range(0_f64..=max_seconds)))
138            .collect::<Vec<_>>();
139        timestamps.sort();
140        for timestamp in timestamps {
141            let sender_id = *zed_user_ids.choose(&mut rng).unwrap();
142            let body = lipsum::lipsum_words(rng.gen_range(1..=50));
143            db.create_channel_message(channel_id, sender_id, &body, timestamp, rng.gen())
144                .await
145                .expect("failed to insert message");
146        }
147        channel_id
148    };
149
150    for user_id in zed_user_ids {
151        db.add_org_member(zed_org_id, user_id, true)
152            .await
153            .expect("failed to insert org membership");
154        db.add_channel_member(general_channel_id, user_id, true)
155            .await
156            .expect("failed to insert channel membership");
157    }
158}
159
160async fn fetch_github<T: DeserializeOwned>(
161    client: &reqwest::Client,
162    access_token: &str,
163    url: &str,
164) -> T {
165    let response = client
166        .get(url)
167        .bearer_auth(&access_token)
168        .header("user-agent", "zed")
169        .send()
170        .await
171        .expect(&format!("failed to fetch '{}'", url));
172    response
173        .json()
174        .await
175        .expect(&format!("failed to deserialize github user from '{}'", url))
176}