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