use anyhow::Result;
use axum::http::StatusCode;
use clap::Parser;
use yorishiro_core::db::TenantDb;
use yorishiro_server::admin::{self, AdminCommand};
use yorishiro_server::{
AppState, bind_addr_from_env, build_app, build_embedding_provider, database_url_from_env,
shutdown_signal,
};
#[derive(Parser)]
#[command(name = "yorishiro-ce-server")]
struct Cli {
#[command(subcommand)]
command: Option<Command>,
}
#[derive(clap::Subcommand)]
enum Command {
Admin {
#[command(subcommand)]
command: AdminCommand,
},
}
fn no_web_ui() -> axum::routing::MethodRouter {
axum::routing::any(|| async { StatusCode::NOT_FOUND })
}
fn main() -> Result<()> {
unsafe {
yorishiro_server::config::load_and_apply_env_overrides()?;
if std::env::var_os("YORISHIRO_MAX_TENANTS").is_none() {
std::env::set_var("YORISHIRO_MAX_TENANTS", "1");
}
}
let cli = Cli::parse();
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?
.block_on(run(cli))
}
async fn run(cli: Cli) -> Result<()> {
let database_url =
database_url_from_env().unwrap_or_else(yorishiro_server::exit_with_config_code);
let identity_pool = sqlx::PgPool::connect(&database_url).await?;
sqlx::migrate!("./migrations").run(&identity_pool).await?;
if let Some(Command::Admin { command }) = cli.command {
return admin::run_with_pool(&identity_pool, command).await;
}
let _log_guard = yorishiro_server::logging::init()?;
tracing::info!("database connected and migrations applied");
let bind_addr = bind_addr_from_env();
let tenant_db = TenantDb::connect(&database_url, 20).await?;
let embedding_provider = build_embedding_provider()?;
let guard_pool = identity_pool.clone();
let state = AppState::new(tenant_db, identity_pool, embedding_provider);
let embedding_tasks = state.embedding_tasks().clone();
let app = build_app(state, no_web_ui());
if let Some(guard) = yorishiro_core::services::db_load_guard::LoadGuardConfig::from_env() {
tracing::info!(
threshold = guard.threshold,
"db load guard enabled: read-only above this many active connections"
);
tokio::spawn(yorishiro_core::services::db_load_guard::run(
guard_pool, guard,
));
}
let listener = tokio::net::TcpListener::bind(&bind_addr).await?;
tracing::info!("yorishiro-ce-server listening on {bind_addr}");
axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.with_graceful_shutdown(shutdown_signal())
.await?;
embedding_tasks.close();
tokio::select! {
result = tokio::time::timeout(std::time::Duration::from_secs(30), embedding_tasks.wait()) => {
if result.is_err() {
tracing::warn!(
"embedding syncs did not finish within 30s; exiting anyway \
(recover with `admin resync-embeddings`)"
);
}
}
_ = shutdown_signal() => {
tracing::warn!(
"second interrupt received; exiting immediately without waiting for embedding \
syncs to finish (recover with `admin resync-embeddings`)"
);
}
}
Ok(())
}