Make it (a tiny bit) easier to run your own collab (#9557)

* Allow creating channels when seeding
* Allow configuring a custom `SEED_PATH`
* Seed the database when creating/migrating it so you don't need a
  separate step for this.

Release Notes:

- N/A

crates/collab/.admins.default.json

@@ -1,8 +0,0 @@
-[
-  "nathansobo",
-  "as-cii",
-  "maxbrunsfeld",
-  "iamnbutler",
-  "mikayla-maki",
-  "JosephTLyons"
-]

crates/collab/.env.toml

@@ -1,4 +1,5 @@
 DATABASE_URL = "postgres://postgres@localhost/zed"
+# DATABASE_URL = "sqlite:////home/zed/.config/zed/db.sqlite3?mode=rwc"
 DATABASE_MAX_CONNECTIONS = 5
 HTTP_PORT = 8080
 API_TOKEN = "secret"
@@ -13,6 +14,7 @@ BLOB_STORE_BUCKET = "the-extensions-bucket"
 BLOB_STORE_URL = "http://127.0.0.1:9000"
 BLOB_STORE_REGION = "the-region"
 ZED_CLIENT_CHECKSUM_SEED = "development-checksum-seed"
+SEED_PATH = "crates/collab/seed.default.json"
 
 # CLICKHOUSE_URL = ""
 # CLICKHOUSE_USER = "default"

crates/collab/Cargo.toml

@@ -13,8 +13,9 @@ workspace = true
 [[bin]]
 name = "collab"
 
-[[bin]]
-name = "seed"
+[features]
+sqlite = ["sea-orm/sqlx-sqlite", "sqlx/sqlite"]
+test-support = ["sqlite"]
 
 [dependencies]
 anyhow.workspace = true

crates/collab/README.md

@@ -6,21 +6,21 @@ It contains our back-end logic for collaboration, to which we connect from the Z
 
 # Local Development
 
- Detailed instructions on getting started are [here](https://zed.dev/docs/local-collaboration).
+Detailed instructions on getting started are [here](https://zed.dev/docs/local-collaboration).
 
 # Deployment
 
 We run two instances of collab:
 
-* Staging (https://staging-collab.zed.dev)
-* Production (https://collab.zed.dev)
+- Staging (https://staging-collab.zed.dev)
+- Production (https://collab.zed.dev)
 
 Both of these run on the Kubernetes cluster hosted in Digital Ocean.
 
 Deployment is triggered by pushing to the `collab-staging` (or `collab-production`) tag in Github. The best way to do this is:
 
-* `./script/deploy-collab staging`
-* `./script/deploy-collab production`
+- `./script/deploy-collab staging`
+- `./script/deploy-collab production`
 
 You can tell what is currently deployed with `./script/what-is-deployed`.
 

crates/collab/seed.default.json

@@ -0,0 +1,12 @@
+{
+  "admins": [
+    "nathansobo",
+    "as-cii",
+    "maxbrunsfeld",
+    "iamnbutler",
+    "mikayla-maki",
+    "JosephTLyons"
+  ],
+  "channels": ["zed"],
+  "number_of_users": 100
+}

crates/collab/src/bin/seed.rs

@@ -1,97 +0,0 @@
-use collab::{
-    db::{self, NewUserParams},
-    env::load_dotenv,
-    executor::Executor,
-};
-use db::{ConnectOptions, Database};
-use serde::{de::DeserializeOwned, Deserialize};
-use std::{fmt::Write, fs};
-
-#[derive(Debug, Deserialize)]
-struct GitHubUser {
-    id: i32,
-    login: String,
-    email: Option<String>,
-}
-
-#[tokio::main]
-async fn main() {
-    load_dotenv().expect("failed to load .env.toml file");
-
-    let mut admin_logins = load_admins("crates/collab/.admins.default.json")
-        .expect("failed to load default admins file");
-    if let Ok(other_admins) = load_admins("./.admins.json") {
-        admin_logins.extend(other_admins);
-    }
-
-    let database_url = std::env::var("DATABASE_URL").expect("missing DATABASE_URL env var");
-    let db = Database::new(ConnectOptions::new(database_url), Executor::Production)
-        .await
-        .expect("failed to connect to postgres database");
-    let client = reqwest::Client::new();
-
-    // Create admin users for all of the users in `.admins.toml` or `.admins.default.toml`.
-    for admin_login in admin_logins {
-        let user = fetch_github::<GitHubUser>(
-            &client,
-            &format!("https://api.github.com/users/{admin_login}"),
-        )
-        .await;
-        db.create_user(
-            &user.email.unwrap_or(format!("{admin_login}@example.com")),
-            true,
-            NewUserParams {
-                github_login: user.login,
-                github_user_id: user.id,
-            },
-        )
-        .await
-        .expect("failed to create admin user");
-    }
-
-    // Fetch 100 other random users from GitHub and insert them into the database.
-    let mut user_count = db
-        .get_all_users(0, 200)
-        .await
-        .expect("failed to load users from db")
-        .len();
-    let mut last_user_id = None;
-    while user_count < 100 {
-        let mut uri = "https://api.github.com/users?per_page=100".to_string();
-        if let Some(last_user_id) = last_user_id {
-            write!(&mut uri, "&since={}", last_user_id).unwrap();
-        }
-        let users = fetch_github::<Vec<GitHubUser>>(&client, &uri).await;
-
-        for github_user in users {
-            last_user_id = Some(github_user.id);
-            user_count += 1;
-            db.get_or_create_user_by_github_account(
-                &github_user.login,
-                Some(github_user.id),
-                github_user.email.as_deref(),
-                None,
-            )
-            .await
-            .expect("failed to insert user");
-        }
-    }
-}
-
-fn load_admins(path: &str) -> anyhow::Result<Vec<String>> {
-    let file_content = fs::read_to_string(path)?;
-    Ok(serde_json::from_str(&file_content)?)
-}
-
-async fn fetch_github<T: DeserializeOwned>(client: &reqwest::Client, url: &str) -> T {
-    let response = client
-        .get(url)
-        .header("user-agent", "zed")
-        .send()
-        .await
-        .unwrap_or_else(|_| panic!("failed to fetch '{}'", url));
-    response
-        .json()
-        .await
-        .unwrap_or_else(|_| panic!("failed to deserialize github user from '{}'", url))
-}

crates/collab/src/db.rs

@@ -128,12 +128,6 @@ impl Database {
         Ok(new_migrations)
     }
 
-    /// Initializes static data that resides in the database by upserting it.
-    pub async fn initialize_static_data(&mut self) -> Result<()> {
-        self.initialize_notification_kinds().await?;
-        Ok(())
-    }
-
     /// Transaction runs things in a transaction. If you want to call other methods
     /// and pass the transaction around you need to reborrow the transaction at each
     /// call site with: `&*tx`.

crates/collab/src/lib.rs

@@ -6,6 +6,7 @@ pub mod env;
 pub mod executor;
 mod rate_limiter;
 pub mod rpc;
+pub mod seed;
 
 #[cfg(test)]
 mod tests;
@@ -111,6 +112,8 @@ impl std::error::Error for Error {}
 pub struct Config {
     pub http_port: u16,
     pub database_url: String,
+    pub migrations_path: Option<PathBuf>,
+    pub seed_path: Option<PathBuf>,
     pub database_max_connections: u32,
     pub api_token: String,
     pub clickhouse_url: Option<String>,
@@ -142,12 +145,6 @@ impl Config {
     }
 }
 
-#[derive(Default, Deserialize)]
-pub struct MigrateConfig {
-    pub database_url: String,
-    pub migrations_path: Option<PathBuf>,
-}
-
 pub struct AppState {
     pub db: Arc<Database>,
     pub live_kit_client: Option<Arc<dyn live_kit_server::api::Client>>,
@@ -162,8 +159,7 @@ impl AppState {
     pub async fn new(config: Config, executor: Executor) -> Result<Arc<Self>> {
         let mut db_options = db::ConnectOptions::new(config.database_url.clone());
         db_options.max_connections(config.database_max_connections);
-        let mut db = Database::new(db_options, Executor::Production).await?;
-        db.initialize_notification_kinds().await?;
+        let db = Database::new(db_options, Executor::Production).await?;
 
         let live_kit_client = if let Some(((server, key), secret)) = config
             .live_kit_server

crates/collab/src/main.rs

@@ -7,7 +7,7 @@ use axum::{
 };
 use collab::{
     api::fetch_extensions_from_blob_store_periodically, db, env, executor::Executor, AppState,
-    Config, MigrateConfig, RateLimiter, Result,
+    Config, RateLimiter, Result,
 };
 use db::Database;
 use std::{
@@ -43,7 +43,16 @@ async fn main() -> Result<()> {
             println!("collab v{} ({})", VERSION, REVISION.unwrap_or("unknown"));
         }
         Some("migrate") => {
-            run_migrations().await?;
+            let config = envy::from_env::<Config>().expect("error loading config");
+            run_migrations(&config).await?;
+        }
+        Some("seed") => {
+            let config = envy::from_env::<Config>().expect("error loading config");
+            let db_options = db::ConnectOptions::new(config.database_url.clone());
+            let mut db = Database::new(db_options, Executor::Production).await?;
+            db.initialize_notification_kinds().await?;
+
+            collab::seed::seed(&config, &db, true).await?;
         }
         Some("serve") => {
             let (is_api, is_collab) = if let Some(next) = args.next() {
@@ -53,14 +62,14 @@ async fn main() -> Result<()> {
             };
             if !is_api && !is_collab {
                 Err(anyhow!(
-                    "usage: collab <version | migrate | serve [api|collab]>"
+                    "usage: collab <version | migrate | seed | serve [api|collab]>"
                 ))?;
             }
 
             let config = envy::from_env::<Config>().expect("error loading config");
             init_tracing(&config);
 
-            run_migrations().await?;
+            run_migrations(&config).await?;
 
             let state = AppState::new(config, Executor::Production).await?;
 
@@ -155,22 +164,25 @@ async fn main() -> Result<()> {
         }
         _ => {
             Err(anyhow!(
-                "usage: collab <version | migrate | serve [api|collab]>"
+                "usage: collab <version | migrate | seed | serve [api|collab]>"
             ))?;
         }
     }
     Ok(())
 }
 
-async fn run_migrations() -> Result<()> {
-    let config = envy::from_env::<MigrateConfig>().expect("error loading config");
+async fn run_migrations(config: &Config) -> Result<()> {
     let db_options = db::ConnectOptions::new(config.database_url.clone());
-    let db = Database::new(db_options, Executor::Production).await?;
+    let mut db = Database::new(db_options, Executor::Production).await?;
 
-    let migrations_path = config
-        .migrations_path
-        .as_deref()
-        .unwrap_or_else(|| Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations")));
+    let migrations_path = config.migrations_path.as_deref().unwrap_or_else(|| {
+        #[cfg(feature = "sqlite")]
+        let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations.sqlite");
+        #[cfg(not(feature = "sqlite"))]
+        let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
+
+        Path::new(default_migrations)
+    });
 
     let migrations = db.migrate(&migrations_path, false).await?;
     for (migration, duration) in migrations {
@@ -182,6 +194,12 @@ async fn run_migrations() -> Result<()> {
         );
     }
 
+    db.initialize_notification_kinds().await?;
+
+    if config.seed_path.is_some() {
+        collab::seed::seed(&config, &db, false).await?;
+    }
+
     return Ok(());
 }
 

crates/collab/src/seed.rs

@@ -0,0 +1,137 @@
+use crate::db::{self, ChannelRole, NewUserParams};
+
+use anyhow::Context;
+use db::Database;
+use serde::{de::DeserializeOwned, Deserialize};
+use std::{fmt::Write, fs, path::Path};
+
+use crate::Config;
+
+#[derive(Debug, Deserialize)]
+struct GitHubUser {
+    id: i32,
+    login: String,
+    email: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct SeedConfig {
+    // Which users to create as admins.
+    admins: Vec<String>,
+    // Which channels to create (all admins are invited to all channels)
+    channels: Vec<String>,
+    // Number of random users to create from the Github API
+    number_of_users: Option<usize>,
+}
+
+pub async fn seed(config: &Config, db: &Database, force: bool) -> anyhow::Result<()> {
+    let client = reqwest::Client::new();
+
+    if !db.get_all_users(0, 1).await?.is_empty() && !force {
+        return Ok(());
+    }
+
+    let seed_path = config
+        .seed_path
+        .as_ref()
+        .context("called seed with no SEED_PATH")?;
+
+    let seed_config = load_admins(seed_path)
+        .context(format!("failed to load {}", seed_path.to_string_lossy()))?;
+
+    let mut first_user = None;
+    let mut others = vec![];
+
+    for admin_login in seed_config.admins {
+        let user = fetch_github::<GitHubUser>(
+            &client,
+            &format!("https://api.github.com/users/{admin_login}"),
+        )
+        .await;
+        let user = db
+            .create_user(
+                &user.email.unwrap_or(format!("{admin_login}@example.com")),
+                true,
+                NewUserParams {
+                    github_login: user.login,
+                    github_user_id: user.id,
+                },
+            )
+            .await
+            .context("failed to create admin user")?;
+        if first_user.is_none() {
+            first_user = Some(user.user_id);
+        } else {
+            others.push(user.user_id)
+        }
+    }
+
+    for channel in seed_config.channels {
+        let (channel, _) = db
+            .create_channel(&channel, None, first_user.unwrap())
+            .await
+            .context("failed to create channel")?;
+
+        for user_id in &others {
+            db.invite_channel_member(
+                channel.id,
+                *user_id,
+                first_user.unwrap(),
+                ChannelRole::Admin,
+            )
+            .await
+            .context("failed to add user to channel")?;
+        }
+    }
+
+    if let Some(number_of_users) = seed_config.number_of_users {
+        // Fetch 100 other random users from GitHub and insert them into the database
+        // (for testing autocompleters, etc.)
+        let mut user_count = db
+            .get_all_users(0, 200)
+            .await
+            .expect("failed to load users from db")
+            .len();
+        let mut last_user_id = None;
+        while user_count < number_of_users {
+            let mut uri = "https://api.github.com/users?per_page=100".to_string();
+            if let Some(last_user_id) = last_user_id {
+                write!(&mut uri, "&since={}", last_user_id).unwrap();
+            }
+            let users = fetch_github::<Vec<GitHubUser>>(&client, &uri).await;
+
+            for github_user in users {
+                last_user_id = Some(github_user.id);
+                user_count += 1;
+                db.get_or_create_user_by_github_account(
+                    &github_user.login,
+                    Some(github_user.id),
+                    github_user.email.as_deref(),
+                    None,
+                )
+                .await
+                .expect("failed to insert user");
+            }
+        }
+    }
+
+    Ok(())
+}
+
+fn load_admins(path: impl AsRef<Path>) -> anyhow::Result<SeedConfig> {
+    let file_content = fs::read_to_string(path)?;
+    Ok(serde_json::from_str(&file_content)?)
+}
+
+async fn fetch_github<T: DeserializeOwned>(client: &reqwest::Client, url: &str) -> T {
+    let response = client
+        .get(url)
+        .header("user-agent", "zed")
+        .send()
+        .await
+        .unwrap_or_else(|_| panic!("failed to fetch '{}'", url));
+    response
+        .json()
+        .await
+        .unwrap_or_else(|_| panic!("failed to deserialize github user from '{}'", url))
+}

crates/collab/src/tests/test_server.rs

@@ -515,6 +515,8 @@ impl TestServer {
                 zed_client_checksum_seed: None,
                 slack_panics_webhook: None,
                 auto_join_channel_id: None,
+                migrations_path: None,
+                seed_path: None,
             },
         })
     }