pg-api 0.1.0

A high-performance PostgreSQL REST API driver with rate limiting, connection pooling, and observability
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

mod models;
mod handlers;
mod auth;
mod database;
mod error;
mod introspection;
mod setup;
mod config;
mod rate_limit;
mod connection_limit;
mod license;
mod observability;

use axum::{
    middleware,
    routing::{get, post, delete},
    Router,
};
use std::net::SocketAddr;
use std::env;
use tower_http::cors::CorsLayer;
use tracing::info;

use crate::models::AppState;
use crate::config::load_server_config;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Load server configuration
    let server_config = load_server_config().await?;
    
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_env_filter(format!("pg_api={}", server_config.log_level))
        .json()
        .init();

    // Check if running in setup mode
    let args: Vec<String> = env::args().collect();
    if args.len() > 1 && args[1] == "setup" {
        return setup::run_setup().await;
    }

    // Initialize app state (includes license)
    let state = AppState::new().await?;
    
    // Initialize observability
    let observability_config = observability::ObservabilityConfig::from_env();
    let observability_client = std::sync::Arc::new(
        observability::ObservabilityClient::new(observability_config.clone())
    );
    
    // Start observability flush task if enabled
    if observability_config.enabled {
        observability_client.clone().start_flush_task();
        info!("Observability enabled, sending metrics to: {:?}", 
              observability_config.opensearch_url);
    } else {
        info!("Observability disabled");
    }

    // Build router
    let app = Router::new()
        // Documentation
        .route("/", get(handlers::serve_docs))
        .route("/docs", get(handlers::serve_docs))
        .route("/openapi.json", get(handlers::serve_openapi))
        
        // Public endpoints
        .route("/health", get(handlers::health_check))
        .route("/v1/status", get(handlers::status_handler))
        .route("/v1/license", get(license::get_license_info))
        
        // Query endpoints
        .route("/v1/query", post(handlers::query_handler))
        .route("/v1/batch", post(handlers::batch_query_handler))
        .route("/v1/transaction", post(handlers::transaction_handler))
        
        // Database management
        .route("/v1/databases", get(handlers::list_databases))
        .route("/v1/databases", post(handlers::create_database))
        .route("/v1/databases/{name}", delete(handlers::drop_database))
        
        // Schema operations
        .route("/v1/databases/{db}/tables", get(handlers::list_tables))
        .route("/v1/databases/{db}/schema", get(handlers::get_schema))
        
        // Account management
        .route("/v1/account", get(handlers::get_account_info))
        .route("/v1/account/usage", get(handlers::get_usage_stats))
        
        .layer(middleware::from_fn_with_state(observability_client.clone(), observability::metrics_middleware))
        .layer(middleware::from_fn_with_state(state.clone(), license::check_license_middleware))
        .layer(middleware::from_fn_with_state(state.clone(), connection_limit::connection_limit_middleware))
        .layer(middleware::from_fn_with_state(state.clone(), rate_limit::rate_limit_middleware))
        .layer(middleware::from_fn_with_state(state.clone(), auth::auth_middleware))
        .layer(middleware::from_fn(auth::request_id_middleware))
        .layer(CorsLayer::permissive())
        .with_state(state);

    // Start server
    let addr = SocketAddr::new(server_config.host, server_config.port);
    let listener = tokio::net::TcpListener::bind(addr).await?;
    
    info!("pg-api running on {}", addr);
    info!("Documentation available at http://{}/docs", addr);
    axum::serve(listener, app).await?;
    
    Ok(())
}