headless_lms_server/programs/seed/
mod.rs

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

headless_lms_server/programs/seed/
mod.rs
