Skip to main content

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_base::config::ApplicationConfiguration;
35use headless_lms_utils::futures::run_parallelly;
36use sqlx::{Pool, Postgres, migrate::MigrateDatabase, postgres::PgPoolOptions};
37use tracing::info;
38
39pub async fn main() -> anyhow::Result<()> {
40    let base_url = ProgramConfig::required("BASE_URL")?;
41    let db_pool = setup_seed_environment().await?;
42    let app_config = ApplicationConfiguration::mock_conf()?;
43    let jwt_password = secrecy::SecretString::new(ProgramConfig::required("JWT_PASSWORD")?.into());
44    let jwt_key = Arc::new(JwtKey::new(&jwt_password).expect("Failed to create JwtKey"));
45
46    // Initialize the global spec fetcher before any seeding
47    seed_helpers::init_seed_spec_fetcher(base_url.clone(), Arc::clone(&jwt_key))
48        .expect("Failed to initialize seed spec fetcher");
49
50    // Run parallelly to improve performance.
51    let (_, seed_users_result, _, seed_llms_result) = try_join!(
52        run_parallelly(seed_exercise_services::seed_exercise_services(
53            db_pool.clone()
54        )),
55        run_parallelly(seed_users::seed_users(db_pool.clone())),
56        run_parallelly(seed_playground_examples::seed_playground_examples(
57            db_pool.clone()
58        )),
59        run_parallelly(seed_application_task_llms(db_pool.clone()))
60    )?;
61
62    // Not run parallely because waits another future that is not send.
63    let seed_file_storage_result = seed_file_storage::seed_file_storage().await?;
64
65    let (uh_cs_organization_result, _uh_mathstat_organization_id, _no_users_organization_id) = try_join!(
66        run_parallelly(seed_organizations::uh_cs::seed_organization_uh_cs(
67            db_pool.clone(),
68            app_config.clone(),
69            seed_users_result,
70            base_url.clone(),
71            Arc::clone(&jwt_key),
72            seed_file_storage_result.clone()
73        )),
74        run_parallelly(
75            seed_organizations::uh_mathstat::seed_organization_uh_mathstat(
76                db_pool.clone(),
77                app_config.clone(),
78                seed_users_result,
79                seed_llms_result,
80                base_url.clone(),
81                Arc::clone(&jwt_key),
82                seed_file_storage_result.clone()
83            )
84        ),
85        run_parallelly(seed_organizations::no_users::seed_organization_no_users(
86            db_pool.clone()
87        ))
88    )?;
89
90    // Sequential rather than in the group above: it shares the default study registry registrar with
91    // the graded course, and get-or-create is not atomic.
92    seed_organizations::credit_registration::seed_organization_credit_registration(
93        db_pool.clone(),
94        &app_config,
95        seed_users_result,
96        base_url.clone(),
97        Arc::clone(&jwt_key),
98        seed_file_storage_result.clone(),
99    )
100    .await?;
101
102    try_join!(
103        run_parallelly(seed_roles::seed_roles(
104            db_pool.clone(),
105            seed_users_result,
106            uh_cs_organization_result
107        )),
108        run_parallelly(seed_user_research_consents::seed_user_research_consents(
109            db_pool.clone(),
110            seed_users_result
111        )),
112        run_parallelly(seed_certificate_fonts::seed_certificate_fonts(
113            db_pool.clone()
114        )),
115        run_parallelly(seed_generic_emails::seed_generic_emails(
116            db_pool.clone(),
117            seed_users_result
118        )),
119        run_parallelly(seed_oauth_clients(db_pool.clone()))
120    )?;
121    seed_global_chatbots::seed_global_chatbots(db_pool.clone()).await?;
122    Ok(())
123}
124
125async fn setup_seed_environment() -> anyhow::Result<Pool<Postgres>> {
126    // TODO: Audit that the environment access only happens in single-threaded code.
127    unsafe { env::set_var("RUST_LOG", "info,sqlx=warn,headless_lms_models=info") };
128
129    dotenvy::dotenv().ok();
130    setup_tracing()?;
131
132    let clean = env::args().any(|a| a == "clean");
133
134    let db_url = ProgramConfig::required("DATABASE_URL")?;
135    let cpu_count = std::thread::available_parallelism()
136        .map(|n| n.get())
137        .unwrap_or(2);
138
139    let max_conns: u32 = std::cmp::max(2, cpu_count as u32);
140
141    let min_conns: u32 = std::cmp::max(1, (cpu_count / 2) as u32);
142
143    let db_pool = PgPoolOptions::new()
144        .max_connections(max_conns)
145        .min_connections(min_conns)
146        // Since this is the seed, it should be fine to wait for a long time for connections
147        .acquire_timeout(Duration::from_secs(10 * 60))
148        .connect(&db_url)
149        .await?;
150
151    if clean {
152        info!("cleaning");
153        // hardcoded for now
154        let status = Command::new("dropdb")
155            .args(["-U", "headless-lms"])
156            .args(["-h", "localhost"])
157            .args(["-p", "54328"])
158            .arg("--force")
159            .arg("-e")
160            .arg("headless_lms_dev")
161            .status()?;
162        assert!(status.success());
163        let db_url = ProgramConfig::required("DATABASE_URL")?;
164        Postgres::create_database(&db_url).await?;
165    }
166
167    if clean {
168        let mut conn = db_pool.acquire().await?;
169        info!("running migrations");
170        sqlx::migrate!("../migrations").run(&mut conn).await?;
171    }
172    Ok(db_pool)
173}