feat: add Config struct loaded from env vars
This commit is contained in:
@@ -0,0 +1,83 @@
|
||||
use std::env;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
pub upstream_url: String,
|
||||
pub upstream_api_key: String,
|
||||
pub database_url: String,
|
||||
pub bind_addr: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn from_env() -> anyhow::Result<Self> {
|
||||
Ok(Self {
|
||||
upstream_url: env::var("UPSTREAM_URL")
|
||||
.unwrap_or_else(|_| "https://api.openai.com".to_string()),
|
||||
upstream_api_key: env::var("UPSTREAM_API_KEY")
|
||||
.map_err(|_| anyhow::anyhow!("UPSTREAM_API_KEY must be set"))?,
|
||||
database_url: env::var("DATABASE_URL")
|
||||
.map_err(|_| anyhow::anyhow!("DATABASE_URL must be set"))?,
|
||||
bind_addr: env::var("BIND_ADDR")
|
||||
.unwrap_or_else(|_| "0.0.0.0".to_string()),
|
||||
port: env::var("PORT")
|
||||
.unwrap_or_else(|_| "8080".to_string())
|
||||
.parse()
|
||||
.map_err(|_| anyhow::anyhow!("PORT must be a valid u16"))?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn listen_addr(&self) -> String {
|
||||
format!("{}:{}", self.bind_addr, self.port)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_from_env_missing_required() {
|
||||
env::remove_var("UPSTREAM_API_KEY");
|
||||
env::remove_var("DATABASE_URL");
|
||||
let result = Config::from_env();
|
||||
assert!(result.is_err(), "Expected error when required vars missing");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_env_with_all_vars() {
|
||||
env::set_var("UPSTREAM_API_KEY", "test-key");
|
||||
env::set_var("DATABASE_URL", "postgres://u:p@localhost/db");
|
||||
env::set_var("UPSTREAM_URL", "https://custom.api.com");
|
||||
env::set_var("PORT", "9090");
|
||||
env::set_var("BIND_ADDR", "127.0.0.1");
|
||||
|
||||
let cfg = Config::from_env().expect("config should load");
|
||||
assert_eq!(cfg.upstream_api_key, "test-key");
|
||||
assert_eq!(cfg.database_url, "postgres://u:p@localhost/db");
|
||||
assert_eq!(cfg.upstream_url, "https://custom.api.com");
|
||||
assert_eq!(cfg.port, 9090);
|
||||
assert_eq!(cfg.bind_addr, "127.0.0.1");
|
||||
assert_eq!(cfg.listen_addr(), "127.0.0.1:9090");
|
||||
|
||||
env::remove_var("UPSTREAM_API_KEY");
|
||||
env::remove_var("DATABASE_URL");
|
||||
env::remove_var("UPSTREAM_URL");
|
||||
env::remove_var("PORT");
|
||||
env::remove_var("BIND_ADDR");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_listen_addr_default_port() {
|
||||
env::set_var("UPSTREAM_API_KEY", "k");
|
||||
env::set_var("DATABASE_URL", "postgres://u:p@localhost/db");
|
||||
env::remove_var("PORT");
|
||||
env::remove_var("BIND_ADDR");
|
||||
|
||||
let cfg = Config::from_env().expect("should load");
|
||||
assert_eq!(cfg.listen_addr(), "0.0.0.0:8080");
|
||||
|
||||
env::remove_var("UPSTREAM_API_KEY");
|
||||
env::remove_var("DATABASE_URL");
|
||||
}
|
||||
}
|
||||
@@ -1 +1,68 @@
|
||||
fn main() { println!("openai-proxy stub"); }
|
||||
mod cache;
|
||||
mod config;
|
||||
mod proxy;
|
||||
mod routes;
|
||||
mod types;
|
||||
|
||||
use anyhow::Result;
|
||||
use axum::{
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
use dotenv::dotenv;
|
||||
use sqlx::postgres::PgPoolOptions;
|
||||
use tower_http::{
|
||||
cors::{Any, CorsLayer},
|
||||
trace::TraceLayer,
|
||||
};
|
||||
use tracing::info;
|
||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
|
||||
|
||||
use crate::cache::CacheStore;
|
||||
use crate::config::Config;
|
||||
use crate::proxy::ProxyClient;
|
||||
use crate::routes::AppState;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
dotenv().ok();
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()))
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
let config = Config::from_env()?;
|
||||
info!("Starting openai-proxy on {}", config.listen_addr());
|
||||
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(10)
|
||||
.connect(&config.database_url)
|
||||
.await?;
|
||||
|
||||
let cache = CacheStore::new(pool);
|
||||
cache.init_schema().await?;
|
||||
info!("Database schema ready");
|
||||
|
||||
let proxy = ProxyClient::new(config.clone(), cache)?;
|
||||
let state = AppState { proxy };
|
||||
|
||||
let cors = CorsLayer::new()
|
||||
.allow_origin(Any)
|
||||
.allow_methods(Any)
|
||||
.allow_headers(Any);
|
||||
|
||||
let app = Router::new()
|
||||
.route("/health", get(routes::health))
|
||||
.route("/v1/chat/completions", post(routes::chat_completions))
|
||||
.route("/v1/models", get(routes::models))
|
||||
.layer(TraceLayer::new_for_http())
|
||||
.layer(cors)
|
||||
.with_state(state);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(config.listen_addr()).await?;
|
||||
info!("Listening on {}", config.listen_addr());
|
||||
axum::serve(listener, app).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user