headless_lms_server/programs/seed/
mod.rs

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

headless_lms_server/programs/seed/
mod.rs
