headless_lms_server/programs/seed/
mod.rs

1#![allow(clippy::unwrap_used)]
2
3pub mod builder;
4pub mod seed_certificate_fonts;
5pub mod seed_courses;
6pub mod seed_exercise_services;
7pub mod seed_file_storage;
8pub mod seed_generic_emails;
9pub mod seed_helpers;
10pub mod seed_oauth_clients;
11pub mod seed_organizations;
12pub mod seed_playground_examples;
13pub mod seed_roles;
14mod seed_user_research_consents;
15pub mod seed_users;
16
17use std::{env, process::Command, sync::Arc, time::Duration};
18
19use crate::{
20    domain::models_requests::JwtKey, programs::seed::seed_oauth_clients::seed_oauth_clients,
21    setup_tracing,
22};
23
24use anyhow::Context;
25use futures::try_join;
26
27use headless_lms_utils::futures::run_parallelly;
28use sqlx::{Pool, Postgres, migrate::MigrateDatabase, postgres::PgPoolOptions};
29use tracing::info;
30
31pub async fn main() -> anyhow::Result<()> {
32    let base_url = std::env::var("BASE_URL").context("BASE_URL must be defined")?;
33    let db_pool = setup_seed_environment().await?;
34    let jwt_key = Arc::new(JwtKey::try_from_env().expect("Failed to create JwtKey"));
35
36    // Initialize the global spec fetcher before any seeding
37    seed_helpers::init_seed_spec_fetcher(base_url.clone(), Arc::clone(&jwt_key))
38        .expect("Failed to initialize seed spec fetcher");
39
40    // Run parallelly to improve performance.
41    let (_, seed_users_result, _) = try_join!(
42        run_parallelly(seed_exercise_services::seed_exercise_services(
43            db_pool.clone()
44        )),
45        run_parallelly(seed_users::seed_users(db_pool.clone())),
46        run_parallelly(seed_playground_examples::seed_playground_examples(
47            db_pool.clone()
48        )),
49    )?;
50
51    // Not run parallely because waits another future that is not send.
52    let seed_file_storage_result = seed_file_storage::seed_file_storage().await?;
53
54    let (uh_cs_organization_result, _uh_mathstat_organization_id, _no_users_organization_id) = try_join!(
55        run_parallelly(seed_organizations::uh_cs::seed_organization_uh_cs(
56            db_pool.clone(),
57            seed_users_result,
58            base_url.clone(),
59            Arc::clone(&jwt_key),
60            seed_file_storage_result.clone()
61        )),
62        run_parallelly(
63            seed_organizations::uh_mathstat::seed_organization_uh_mathstat(
64                db_pool.clone(),
65                seed_users_result,
66                base_url.clone(),
67                Arc::clone(&jwt_key),
68                seed_file_storage_result.clone()
69            )
70        ),
71        run_parallelly(seed_organizations::no_users::seed_organization_no_users(
72            db_pool.clone()
73        ))
74    )?;
75
76    try_join!(
77        run_parallelly(seed_roles::seed_roles(
78            db_pool.clone(),
79            seed_users_result,
80            uh_cs_organization_result
81        )),
82        run_parallelly(seed_user_research_consents::seed_user_research_consents(
83            db_pool.clone(),
84            seed_users_result
85        )),
86        run_parallelly(seed_certificate_fonts::seed_certificate_fonts(
87            db_pool.clone()
88        )),
89        run_parallelly(seed_generic_emails::seed_generic_emails(
90            db_pool.clone(),
91            seed_users_result
92        )),
93        run_parallelly(seed_oauth_clients(db_pool.clone()))
94    )?;
95
96    Ok(())
97}
98
99async fn setup_seed_environment() -> anyhow::Result<Pool<Postgres>> {
100    // TODO: Audit that the environment access only happens in single-threaded code.
101    unsafe { env::set_var("RUST_LOG", "info,sqlx=warn,headless_lms_models=info") };
102
103    dotenv::dotenv().ok();
104    setup_tracing()?;
105
106    let clean = env::args().any(|a| a == "clean");
107
108    let db_url = env::var("DATABASE_URL")?;
109    let cpu_count = std::thread::available_parallelism()
110        .map(|n| n.get())
111        .unwrap_or(2);
112
113    let max_conns: u32 = std::cmp::max(2, cpu_count as u32);
114
115    let min_conns: u32 = std::cmp::max(1, (cpu_count / 2) as u32);
116
117    let db_pool = PgPoolOptions::new()
118        .max_connections(max_conns)
119        .min_connections(min_conns)
120        // Since this is the seed, it should be fine to wait for a long time for connections
121        .acquire_timeout(Duration::from_secs(10 * 60))
122        .connect(&db_url)
123        .await?;
124
125    if clean {
126        info!("cleaning");
127        // hardcoded for now
128        let status = Command::new("dropdb")
129            .args(["-U", "headless-lms"])
130            .args(["-h", "localhost"])
131            .args(["-p", "54328"])
132            .arg("--force")
133            .arg("-e")
134            .arg("headless_lms_dev")
135            .status()?;
136        assert!(status.success());
137        let db_url = env::var("DATABASE_URL")?;
138        Postgres::create_database(&db_url).await?;
139    }
140
141    if clean {
142        let mut conn = db_pool.acquire().await?;
143        info!("running migrations");
144        sqlx::migrate!("../migrations").run(&mut conn).await?;
145    }
146    Ok(db_pool)
147}