diff --git a/rust/rsc/src/bin/rsc/config.rs b/rust/rsc/src/bin/rsc/config.rs new file mode 100644 index 000000000..1e12491bb --- /dev/null +++ b/rust/rsc/src/bin/rsc/config.rs @@ -0,0 +1,21 @@ +use config::{Config, ConfigError, Environment, File}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Serialize)] +pub struct RSCConfig { + pub database_url: String, + pub server_addr: String, + pub active_store: String, +} + +impl RSCConfig { + pub fn new() -> Result { + // Gather the config + let config = Config::builder() + .add_source(Environment::with_prefix("WAKE_RSC_CONFIG")) + .add_source(File::with_name(".config")) + .build()?; + + config.try_deserialize() + } +} diff --git a/rust/rsc/src/bin/rsc/main.rs b/rust/rsc/src/bin/rsc/main.rs index 0fafd3d2d..37f7ef7a6 100644 --- a/rust/rsc/src/bin/rsc/main.rs +++ b/rust/rsc/src/bin/rsc/main.rs @@ -5,15 +5,10 @@ use axum::{ }; use chrono::Utc; use clap::Parser; -use data_encoding::HEXLOWER; use migration::{Migrator, MigratorTrait}; -use rand_core::{OsRng, RngCore}; use rlimit::Resource; -use rsc::{config, database}; -use sea_orm::{ - prelude::Uuid, ActiveModelTrait, ActiveValue::*, ConnectOptions, ConnectionTrait, Database, - DatabaseConnection, -}; +use rsc::database; +use sea_orm::{prelude::Uuid, ConnectOptions, Database, DatabaseConnection}; use std::collections::HashMap; use std::io::{Error, ErrorKind}; use std::sync::Arc; @@ -24,34 +19,15 @@ mod add_job; mod api_key_check; mod blob; mod blob_store_impls; +mod config; mod dashboard; mod read_job; mod types; #[derive(Debug, Parser)] struct ServerOptions { - #[arg(help = "Specify a config override file", value_name = "CONFIG", long)] - config_override: Option, - #[arg(help = "Shows the config and then exits", long)] show_config: bool, - - #[arg( - help = "Specify an override for the bind address", - value_name = "SERVER_IP[:SERVER_PORT]", - long - )] - server_addr: Option, - - #[arg( - help = "Specify an override for the database url", - value_name = "DATABASE_URL", - long - )] - database_url: Option, - - #[arg(help = "Launches the cache without an external database", long)] - standalone: bool, } // Maps databse store uuid -> dyn blob::DebugBlobStore @@ -134,12 +110,7 @@ fn create_router( config: Arc, blob_stores: &HashMap>, ) -> Router { - // If we can't create a store, just exit. The config is wrong and must be rectified. - let Some(active_store_uuid) = config.active_store.clone() else { - panic!("Active store uuid not set in configuration"); - }; - - let Ok(active_store_uuid) = Uuid::parse_str(&active_store_uuid) else { + let Ok(active_store_uuid) = Uuid::parse_str(&config.active_store) else { panic!("Failed to parse provided active store into uuid"); }; @@ -215,22 +186,7 @@ fn create_router( .route("/version/check", get(check_version)) } -async fn create_standalone_db() -> Result { - let shim_db = Database::connect("postgres://127.0.0.1/shim").await?; - let mut buf = [0u8; 24]; - OsRng.fill_bytes(&mut buf); - let rng_str = HEXLOWER.encode(&buf); - let db = format!("db_{}", rng_str); - shim_db - .execute_unprepared(&format!("CREATE DATABASE {}", db)) - .await?; - drop(shim_db); - let db = Database::connect(format!("postgres://127.0.0.1/{}", db)).await?; - Migrator::up(&db, None).await?; - Ok(db) -} - -async fn create_remote_db( +async fn connect_to_database( config: &config::RSCConfig, ) -> Result> { let mut opt = ConnectOptions::new(&config.database_url); @@ -252,36 +208,6 @@ async fn create_remote_db( Ok(connection) } -async fn create_insecure_api_key( - db: &DatabaseConnection, -) -> Result> { - let active_key = entity::api_key::ActiveModel { - id: NotSet, - created_at: NotSet, - key: Set("InsecureKey".into()), - desc: Set("Generated Insecure Key".into()), - }; - - let inserted_key = active_key.insert(db).await?; - - Ok(inserted_key.key) -} - -async fn connect_to_database( - config: &config::RSCConfig, -) -> Result> { - if config.standalone { - tracing::warn!("Launching rsc in standalone mode, data will not persist."); - let db = create_standalone_db().await?; - let key = create_insecure_api_key(&db).await?; - tracing::info!(key, "Created insecure api key."); - - return Ok(db); - } - - create_remote_db(config).await -} - fn launch_job_eviction(conn: Arc, tick_interval: u64, ttl: u64) { tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_secs(tick_interval)); @@ -393,17 +319,7 @@ async fn main() -> Result<(), Box> { let args = ServerOptions::parse(); // Get the configuration - let config = config::RSCConfig::new(config::RSCConfigOverride { - config_override: args.config_override, - server_addr: args.server_addr, - database_url: args.database_url, - standalone: if args.standalone { - Some(args.standalone) - } else { - None - }, - active_store: None, - })?; + let config = config::RSCConfig::new()?; let config = Arc::new(config); if args.show_config { @@ -446,8 +362,13 @@ async fn main() -> Result<(), Box> { #[cfg(test)] mod tests { use super::*; + use data_encoding::HEXLOWER; use entity::blob_store; - use sea_orm::{prelude::Uuid, ActiveModelTrait, EntityTrait, PaginatorTrait}; + use rand_core::{OsRng, RngCore}; + use sea_orm::{ + prelude::Uuid, ActiveModelTrait, ActiveValue::*, ConnectionTrait, EntityTrait, + PaginatorTrait, + }; use std::sync::Arc; use axum::{ @@ -474,14 +395,12 @@ mod tests { Ok(inserted.id) } - fn create_config(store_id: Uuid) -> Result> { - Ok(config::RSCConfig::new(config::RSCConfigOverride { - config_override: Some("".into()), - server_addr: Some("test:0000".into()), - database_url: Some("".into()), - standalone: Some(true), - active_store: Some(store_id.to_string()), - })?) + fn create_config(store_id: Uuid) -> config::RSCConfig { + config::RSCConfig { + database_url: "test:0000".to_string(), + server_addr: "".to_string(), + active_store: store_id.to_string(), + } } async fn create_fake_blob( @@ -502,13 +421,43 @@ mod tests { Ok(inserted_key.id) } + async fn create_insecure_api_key( + db: &DatabaseConnection, + ) -> Result> { + let active_key = entity::api_key::ActiveModel { + id: NotSet, + created_at: NotSet, + key: Set("InsecureKey".into()), + desc: Set("Generated Insecure Key".into()), + }; + + let inserted_key = active_key.insert(db).await?; + + Ok(inserted_key.key) + } + + async fn create_standalone_db() -> Result { + let shim_db = Database::connect("postgres://127.0.0.1/shim").await?; + let mut buf = [0u8; 24]; + OsRng.fill_bytes(&mut buf); + let rng_str = HEXLOWER.encode(&buf); + let db = format!("db_{}", rng_str); + shim_db + .execute_unprepared(&format!("CREATE DATABASE {}", db)) + .await?; + drop(shim_db); + let db = Database::connect(format!("postgres://127.0.0.1/{}", db)).await?; + Migrator::up(&db, None).await?; + Ok(db) + } + #[tokio::test] async fn nominal() { let db = create_standalone_db().await.unwrap(); let store_id = create_test_store(&db).await.unwrap(); let api_key = create_insecure_api_key(&db).await.unwrap(); let blob_id = create_fake_blob(&db, store_id.clone()).await.unwrap(); - let config = create_config(store_id.clone()).unwrap(); + let config = create_config(store_id.clone()); let db = Arc::new(db); let stores = activate_stores(db.clone()).await; let mut router = create_router(db.clone(), Arc::new(config), &stores); diff --git a/rust/rsc/src/bin/rsc_tool/config.rs b/rust/rsc/src/bin/rsc_tool/config.rs new file mode 100644 index 000000000..2e88e9396 --- /dev/null +++ b/rust/rsc/src/bin/rsc_tool/config.rs @@ -0,0 +1,25 @@ +use config::{Config, ConfigError, Environment, File}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Deserialize, Default)] +pub struct RSCToolConfigOverride { + pub database_url: Option, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct RSCToolConfig { + pub database_url: String, +} + +impl RSCToolConfig { + pub fn new(overrides: RSCToolConfigOverride) -> Result { + // Gather the config + let config = Config::builder() + .add_source(Environment::with_prefix("WAKE_RSC_CONFIG")) + .add_source(File::with_name(".config").required(false)) + .set_override_option("database_url", overrides.database_url)? + .build()?; + + config.try_deserialize() + } +} diff --git a/rust/rsc/src/bin/rsc_tool/main.rs b/rust/rsc/src/bin/rsc_tool/main.rs index 59dc6e538..397309db0 100644 --- a/rust/rsc/src/bin/rsc_tool/main.rs +++ b/rust/rsc/src/bin/rsc_tool/main.rs @@ -2,11 +2,12 @@ use clap::{Parser, Subcommand}; use inquire::{Confirm, Text}; use is_terminal::IsTerminal; use migration::{DbErr, Migrator, MigratorTrait}; -use rsc::{config, database}; +use rsc::database; use sea_orm::{prelude::Uuid, DatabaseConnection}; use std::io::{Error, ErrorKind}; use tracing; +mod config; mod table; async fn add_api_key( @@ -242,14 +243,6 @@ async fn bootstrap_db(db: &DatabaseConnection) -> Result<(), Box, - #[arg( help = "Specify and override for the database url", value_name = "DATABASE_URL", @@ -396,10 +389,8 @@ async fn main() -> Result<(), Box> { let args = TopLevel::parse(); // Gather our config - let config = config::RSCConfig::new(config::RSCConfigOverride { - config_override: args.config_override, + let config = config::RSCToolConfig::new(config::RSCToolConfigOverride { database_url: args.database_url, - ..Default::default() })?; if args.show_config { diff --git a/rust/rsc/src/config.rs b/rust/rsc/src/config.rs deleted file mode 100644 index cdbbd1fc2..000000000 --- a/rust/rsc/src/config.rs +++ /dev/null @@ -1,47 +0,0 @@ -use config::{Config, ConfigError, Environment, File}; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Deserialize, Default)] -pub struct RSCConfigOverride { - pub config_override: Option, - pub database_url: Option, - pub server_addr: Option, - pub standalone: Option, - // Active store is specified by setting the UUID of the store - // that should be processing uploads. - pub active_store: Option, -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct RSCConfig { - pub database_url: String, - // TODO: We should allow setting a domain as well - pub server_addr: String, - pub standalone: bool, - pub active_store: Option, -} - -impl RSCConfig { - pub fn new(overrides: RSCConfigOverride) -> Result { - // Gather the config - let config = Config::builder() - .add_source(Environment::with_prefix("WAKE_RSC")) - .add_source( - File::with_name( - &overrides - .config_override - .as_ref() - .map(|x| x.as_str()) - .unwrap_or(".config"), - ) - .required(false), - ) - .set_override_option("database_url", overrides.database_url)? - .set_override_option("server_addr", overrides.server_addr)? - .set_override_option("standalone", overrides.standalone)? - .set_override_option("active_store", overrides.active_store)? - .build()?; - - config.try_deserialize() - } -} diff --git a/rust/rsc/src/lib.rs b/rust/rsc/src/lib.rs index 2ee385a05..8fd0a6be8 100644 --- a/rust/rsc/src/lib.rs +++ b/rust/rsc/src/lib.rs @@ -1,2 +1 @@ -pub mod config; pub mod database;