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