postrust_server/
app.rs

1//! Request handling.
2
3use crate::state::AppState;
4use axum::{
5    body::Body,
6    extract::{Request, State},
7    http::StatusCode,
8    response::{IntoResponse, Response},
9};
10use postrust_auth::authenticate;
11use postrust_core::{
12    create_action_plan, parse_request, ActionPlan, ApiRequest, CallPlan, DbActionPlan,
13};
14use postrust_response::{format_response, QueryResult, Response as PgrstResponse};
15use std::sync::Arc;
16use tracing::{debug, error};
17
18/// Main request handler.
19pub async fn handle_request(State(state): State<Arc<AppState>>, request: Request) -> Response {
20    let method = request.method().clone();
21    let path = request.uri().path().to_string();
22
23    debug!("{} {}", method, path);
24
25    match process_request(state, request).await {
26        Ok(response) => response.into_response(),
27        Err(e) => error_response(e).into_response(),
28    }
29}
30
31/// Process a request and return a response.
32async fn process_request(
33    state: Arc<AppState>,
34    request: Request,
35) -> Result<Response, postrust_core::Error> {
36    // Extract auth header
37    let auth_header = request
38        .headers()
39        .get("authorization")
40        .and_then(|v| v.to_str().ok());
41
42    // Authenticate
43    let auth_result = authenticate(auth_header, &state.jwt_config)
44        .map_err(|e| postrust_core::Error::InvalidJwt(e.to_string()))?;
45
46    debug!("Authenticated as role: {}", auth_result.role);
47
48    // Parse request
49    let (parts, body) = request.into_parts();
50    let body_bytes = axum::body::to_bytes(body, 10 * 1024 * 1024)
51        .await
52        .map_err(|e| postrust_core::Error::InvalidBody(e.to_string()))?;
53
54    // Build HTTP request for parsing
55    let mut builder = http::Request::builder()
56        .method(parts.method.clone())
57        .uri(parts.uri.clone());
58
59    for (key, value) in &parts.headers {
60        builder = builder.header(key, value);
61    }
62
63    let http_request = builder
64        .body(body_bytes.clone())
65        .map_err(|e| postrust_core::Error::Internal(e.to_string()))?;
66
67    // Parse API request
68    let mut api_request = parse_request(&http_request, state.default_schema(), state.schemas())?;
69
70    // Parse payload
71    if !body_bytes.is_empty() {
72        let payload = postrust_core::api_request::payload::parse_payload(
73            body_bytes,
74            &api_request.content_media_type,
75        )?;
76        api_request.payload = payload;
77    }
78
79    // Get schema cache
80    let schema_cache = state.schema_cache().await;
81
82    // Create execution plan
83    let plan = create_action_plan(&api_request, &schema_cache)?;
84
85    // Execute plan
86    let result = execute_plan(&state, &api_request, &plan, &auth_result).await?;
87
88    // Format response
89    let response = format_response(&api_request, &result)
90        .map_err(|e| postrust_core::Error::Internal(e.to_string()))?;
91
92    Ok(build_response(response))
93}
94
95/// Execute an action plan.
96async fn execute_plan(
97    state: &AppState,
98    _request: &ApiRequest,
99    plan: &ActionPlan,
100    auth: &postrust_auth::AuthResult,
101) -> Result<QueryResult, postrust_core::Error> {
102    match plan {
103        ActionPlan::Db(db_plan) => {
104            // Build SQL
105            let query = postrust_core::query::build_query(
106                &ActionPlan::Db(db_plan.clone()),
107                Some(&auth.role),
108            )?;
109
110            if !query.has_main() {
111                return Ok(QueryResult::default());
112            }
113
114            let (sql, params) = query.build_main();
115            debug!("Executing SQL: {}", sql);
116            debug!("With {} parameters", params.len());
117
118            // Execute query
119            let mut conn = state
120                .pool
121                .acquire()
122                .await
123                .map_err(|e| postrust_core::Error::ConnectionPool(e.to_string()))?;
124
125            // Set role
126            sqlx::query(&format!(
127                "SET LOCAL ROLE {}",
128                postrust_sql::escape_ident(&auth.role)
129            ))
130            .execute(&mut *conn)
131            .await
132            .map_err(|e| {
133                postrust_core::Error::Database(postrust_core::error::DatabaseError {
134                    code: "42501".into(),
135                    message: e.to_string(),
136                    details: None,
137                    hint: None,
138                    constraint: None,
139                    table: None,
140                    column: None,
141                })
142            })?;
143
144            // Set claims as GUC
145            for (key, value) in &auth.claims {
146                let guc_key = format!("request.jwt.claims.{}", key);
147                let guc_value = match value {
148                    serde_json::Value::String(s) => s.clone(),
149                    other => other.to_string(),
150                };
151
152                sqlx::query("SELECT set_config($1, $2, true)")
153                    .bind(&guc_key)
154                    .bind(&guc_value)
155                    .execute(&mut *conn)
156                    .await
157                    .ok(); // Ignore errors for individual claims
158            }
159
160            // Execute main query with bound parameters
161            let rows = bind_params(sqlx::query(&sql), &params)
162                .fetch_all(&mut *conn)
163                .await
164                .map_err(|e| {
165                    error!("Query error: {}", e);
166                    map_sqlx_error(e)
167                })?;
168
169            // Convert rows to JSON
170            let json_rows: Vec<serde_json::Value> = rows.iter().map(row_to_json).collect();
171
172            // In PostgREST-compatibility mode, reshape RPC responses to match
173            // PostgREST: un-nest the function-name-keyed column and return a
174            // bare value for non-set-returning functions.
175            let (json_rows, singular) = if state.config.compat_mode {
176                if let ActionPlan::Db(DbActionPlan::Call { call, .. }) = plan {
177                    unwrap_rpc_rows(json_rows, call)
178                } else {
179                    (json_rows, false)
180                }
181            } else {
182                (json_rows, false)
183            };
184
185            Ok(QueryResult {
186                status: StatusCode::OK,
187                rows: json_rows,
188                singular,
189                ..Default::default()
190            })
191        }
192        ActionPlan::Info(info_plan) => {
193            use postrust_core::plan::InfoPlan;
194
195            // Return appropriate metadata based on the info type
196            let response_data = match info_plan {
197                InfoPlan::OpenApiSpec => {
198                    // Return basic server info for root endpoint
199                    serde_json::json!({
200                        "name": "postrust",
201                        "version": env!("CARGO_PKG_VERSION"),
202                        "description": "PostgREST-compatible REST API for PostgreSQL"
203                    })
204                }
205                InfoPlan::RelationInfo(qi) => {
206                    serde_json::json!({
207                        "schema": qi.schema,
208                        "name": qi.name,
209                        "type": "relation"
210                    })
211                }
212                InfoPlan::RoutineInfo(qi) => {
213                    serde_json::json!({
214                        "schema": qi.schema,
215                        "name": qi.name,
216                        "type": "routine"
217                    })
218                }
219            };
220
221            Ok(QueryResult {
222                status: StatusCode::OK,
223                rows: vec![response_data],
224                ..Default::default()
225            })
226        }
227    }
228}
229
230/// Reshape RPC rows for PostgREST-compatibility mode.
231///
232/// `SELECT * FROM func(...)` names its single output column after the function
233/// (for scalar/`json` returns), which serializes to rows like
234/// `[{"func": <value>}]`. PostgREST instead returns the bare value. This
235/// un-nests that wrapper column and reports whether the result should be
236/// rendered as a single (un-arrayed) value — i.e. when the function is not
237/// set-returning.
238///
239/// The decision is driven by the plan's return-type metadata: composite and
240/// `record` returns (`RETURNS TABLE`, row types) have real output columns and
241/// are never un-nested, even when a single column happens to share the
242/// function's name (e.g. `CREATE FUNCTION foo() RETURNS TABLE(foo int)`).
243fn unwrap_rpc_rows(
244    rows: Vec<serde_json::Value>,
245    call: &CallPlan,
246) -> (Vec<serde_json::Value>, bool) {
247    let singular = !call.returns_set;
248
249    if call.returns_composite {
250        return (rows, singular);
251    }
252
253    let fname = call.function.name.as_str();
254    let unwrapped = rows
255        .into_iter()
256        .map(|row| match row {
257            serde_json::Value::Object(ref map) if map.len() == 1 && map.contains_key(fname) => {
258                map.get(fname).cloned().unwrap_or(serde_json::Value::Null)
259            }
260            other => other,
261        })
262        .collect();
263
264    (unwrapped, singular)
265}
266
267/// Convert a sqlx row to JSON.
268fn row_to_json(row: &sqlx::postgres::PgRow) -> serde_json::Value {
269    use sqlx::{Column, Row, TypeInfo};
270
271    let mut map = serde_json::Map::new();
272
273    for column in row.columns() {
274        let name = column.name();
275        let type_name = column.type_info().name();
276
277        let value = match type_name {
278            "INT2" | "SMALLINT" => row
279                .try_get::<i16, _>(name)
280                .ok()
281                .map(|v| serde_json::Value::Number(v.into())),
282            "INT4" | "INT" | "INTEGER" => row
283                .try_get::<i32, _>(name)
284                .ok()
285                .map(|v| serde_json::Value::Number(v.into())),
286            "INT8" | "BIGINT" => row
287                .try_get::<i64, _>(name)
288                .ok()
289                .map(|v| serde_json::Value::Number(v.into())),
290            "FLOAT4" | "REAL" => row
291                .try_get::<f32, _>(name)
292                .ok()
293                .and_then(|v| serde_json::Number::from_f64(v as f64))
294                .map(serde_json::Value::Number),
295            "FLOAT8" | "DOUBLE PRECISION" => row
296                .try_get::<f64, _>(name)
297                .ok()
298                .and_then(serde_json::Number::from_f64)
299                .map(serde_json::Value::Number),
300            "NUMERIC" | "DECIMAL" => row
301                .try_get::<sqlx::types::BigDecimal, _>(name)
302                .ok()
303                .map(|v| serde_json::Value::String(v.to_string())),
304            "BOOL" | "BOOLEAN" => row
305                .try_get::<bool, _>(name)
306                .ok()
307                .map(serde_json::Value::Bool),
308            "JSON" | "JSONB" => row.try_get::<serde_json::Value, _>(name).ok(),
309            "UUID" => row
310                .try_get::<sqlx::types::Uuid, _>(name)
311                .ok()
312                .map(|v| serde_json::Value::String(v.to_string())),
313            "TIMESTAMPTZ" | "TIMESTAMP WITH TIME ZONE" => row
314                .try_get::<chrono::DateTime<chrono::Utc>, _>(name)
315                .ok()
316                .map(|v| serde_json::Value::String(v.to_rfc3339())),
317            "TIMESTAMP" | "TIMESTAMP WITHOUT TIME ZONE" => row
318                .try_get::<chrono::NaiveDateTime, _>(name)
319                .ok()
320                .map(|v| serde_json::Value::String(v.to_string())),
321            "DATE" => row
322                .try_get::<chrono::NaiveDate, _>(name)
323                .ok()
324                .map(|v| serde_json::Value::String(v.to_string())),
325            "TIME" | "TIME WITHOUT TIME ZONE" => row
326                .try_get::<chrono::NaiveTime, _>(name)
327                .ok()
328                .map(|v| serde_json::Value::String(v.to_string())),
329            _ => row
330                .try_get::<String, _>(name)
331                .ok()
332                .map(serde_json::Value::String),
333        };
334
335        map.insert(name.to_string(), value.unwrap_or(serde_json::Value::Null));
336    }
337
338    serde_json::Value::Object(map)
339}
340
341/// Bind SqlParam values to a sqlx query.
342fn bind_params<'q>(
343    mut query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
344    params: &'q [postrust_sql::SqlParam],
345) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
346    use postrust_sql::SqlParam;
347
348    for param in params {
349        query = match param {
350            SqlParam::Null => query.bind(None::<String>),
351            SqlParam::Bool(b) => query.bind(b),
352            SqlParam::Int(n) => query.bind(n),
353            SqlParam::Float(f) => query.bind(f),
354            SqlParam::Text(s) => query.bind(s),
355            SqlParam::Bytes(b) => query.bind(b),
356            SqlParam::Json(j) => query.bind(j),
357            SqlParam::Uuid(u) => query.bind(u),
358            SqlParam::Timestamp(t) => query.bind(t),
359            SqlParam::Array(arr) => {
360                // Convert array to Vec<String> for text arrays
361                let strings: Vec<String> = arr
362                    .iter()
363                    .map(|p| match p {
364                        SqlParam::Text(s) => s.clone(),
365                        SqlParam::Int(n) => n.to_string(),
366                        SqlParam::Bool(b) => b.to_string(),
367                        other => format!("{:?}", other),
368                    })
369                    .collect();
370                query.bind(strings)
371            }
372        };
373    }
374
375    query
376}
377
378/// Map sqlx error to our error type.
379fn map_sqlx_error(e: sqlx::Error) -> postrust_core::Error {
380    match e {
381        sqlx::Error::Database(db_err) => {
382            // Try to downcast to Postgres-specific error for additional details
383            let (details, hint) = db_err
384                .try_downcast_ref::<sqlx::postgres::PgDatabaseError>()
385                .map(|pg_err| {
386                    (
387                        pg_err.detail().map(String::from),
388                        pg_err.hint().map(String::from),
389                    )
390                })
391                .unwrap_or((None, None));
392
393            postrust_core::Error::Database(postrust_core::error::DatabaseError {
394                code: db_err.code().map(|c| c.to_string()).unwrap_or_default(),
395                message: db_err.message().to_string(),
396                details,
397                hint,
398                constraint: db_err.constraint().map(|s| s.to_string()),
399                table: db_err.table().map(|s| s.to_string()),
400                column: None,
401            })
402        }
403        other => postrust_core::Error::Internal(other.to_string()),
404    }
405}
406
407/// Build an HTTP response from our response type.
408fn build_response(response: PgrstResponse) -> Response {
409    let mut builder = Response::builder().status(response.status);
410
411    for (key, value) in &response.headers {
412        builder = builder.header(key, value);
413    }
414
415    builder
416        .body(Body::from(response.body))
417        .unwrap_or_else(|_| Response::new(Body::empty()))
418}
419
420/// Build an error response.
421///
422/// In production mode (PGRST_DEBUG=false or unset), sensitive error details
423/// are hidden to prevent information leakage.
424fn error_response(error: postrust_core::Error) -> Response {
425    let status = error.status_code();
426
427    // Check if debug mode is enabled
428    let debug_mode = std::env::var("PGRST_DEBUG")
429        .map(|v| v == "true" || v == "1")
430        .unwrap_or(false);
431
432    let body = if debug_mode {
433        // Full error details in debug mode
434        serde_json::to_vec(&error.to_json()).unwrap_or_default()
435    } else {
436        // Sanitized error in production
437        let sanitized = serde_json::json!({
438            "code": error.code(),
439            "message": sanitize_error_message(&error),
440            "details": null,
441            "hint": null
442        });
443        serde_json::to_vec(&sanitized).unwrap_or_default()
444    };
445
446    Response::builder()
447        .status(status)
448        .header("content-type", "application/json")
449        .body(Body::from(body))
450        .unwrap_or_else(|_| Response::new(Body::empty()))
451}
452
453/// Sanitize error messages for production.
454fn sanitize_error_message(error: &postrust_core::Error) -> &'static str {
455    use postrust_core::Error;
456    match error {
457        Error::TableNotFound(_) | Error::NotFound(_) => "Resource not found",
458        Error::FunctionNotFound(_) => "Function not found",
459        Error::ColumnNotFound(_) | Error::UnknownColumn(_) => "Column not found",
460        Error::RelationshipNotFound(_) => "Relationship not found",
461        Error::InvalidPath(_) => "Invalid request path",
462        Error::InvalidBody(_) => "Invalid request body",
463        Error::InvalidJwt(_) | Error::JwtExpired | Error::MissingAuth => "Unauthorized",
464        Error::InsufficientPermissions(_) => "Forbidden",
465        Error::UnacceptableSchema(_) => "Invalid schema",
466        Error::InvalidHeader(_) | Error::InvalidQueryParam(_) => "Invalid request",
467        Error::Database(_) => "Database error",
468        Error::ConnectionPool(_) => "Service temporarily unavailable",
469        Error::Internal(_) => "Internal server error",
470        _ => "An error occurred",
471    }
472}
473
474#[cfg(test)]
475mod tests {
476    use super::*;
477    use postrust_core::plan::CallParams;
478    use postrust_core::QualifiedIdentifier;
479    use serde_json::json;
480
481    fn call_plan(name: &str, returns_set: bool) -> CallPlan {
482        CallPlan {
483            function: QualifiedIdentifier::new("public", name),
484            params: CallParams::None,
485            returns_scalar: !returns_set,
486            returns_set,
487            returns_composite: false,
488            volatility: "Volatile".into(),
489        }
490    }
491
492    fn composite_call_plan(name: &str, returns_set: bool) -> CallPlan {
493        CallPlan {
494            returns_composite: true,
495            returns_scalar: false,
496            ..call_plan(name, returns_set)
497        }
498    }
499
500    #[test]
501    fn unwraps_json_return_to_bare_object() {
502        // `SELECT * FROM sync(...)` on a json-returning function yields a single
503        // column named after the function.
504        let rows = vec![json!({"sync": {"ok": true, "count": 3}})];
505        let (rows, singular) = unwrap_rpc_rows(rows, &call_plan("sync", false));
506        assert!(singular, "non-set-returning function should be singular");
507        assert_eq!(rows, vec![json!({"ok": true, "count": 3})]);
508    }
509
510    #[test]
511    fn unwraps_scalar_return() {
512        let rows = vec![json!({"add": 42})];
513        let (rows, singular) = unwrap_rpc_rows(rows, &call_plan("add", false));
514        assert!(singular);
515        assert_eq!(rows, vec![json!(42)]);
516    }
517
518    #[test]
519    fn unwraps_setof_scalar_to_array() {
520        let rows = vec![json!({"gen": 1}), json!({"gen": 2})];
521        let (rows, singular) = unwrap_rpc_rows(rows, &call_plan("gen", true));
522        assert!(!singular, "set-returning function should not be singular");
523        assert_eq!(rows, vec![json!(1), json!(2)]);
524    }
525
526    #[test]
527    fn leaves_multi_column_rows_untouched() {
528        // `RETURNS TABLE(...)` / composite set: rows already have real columns
529        // and must not be un-nested.
530        let rows = vec![json!({"id": 1, "name": "a"}), json!({"id": 2, "name": "b"})];
531        let (out, singular) =
532            unwrap_rpc_rows(rows.clone(), &composite_call_plan("list_users", true));
533        assert!(!singular);
534        assert_eq!(out, rows);
535    }
536
537    #[test]
538    fn leaves_table_column_named_like_function_untouched() {
539        // `CREATE FUNCTION foo() RETURNS TABLE(foo int)`: the single output
540        // column legitimately shares the function's name. The composite
541        // return-type metadata must prevent it from being mistaken for the
542        // function-name wrapper.
543        let rows = vec![json!({"foo": 1}), json!({"foo": 2})];
544        let (out, singular) = unwrap_rpc_rows(rows.clone(), &composite_call_plan("foo", true));
545        assert!(!singular);
546        assert_eq!(out, rows);
547    }
548
549    #[test]
550    fn single_composite_return_is_singular_but_not_unwrapped() {
551        // A non-set function returning a row type expands to its columns;
552        // nothing to un-nest, but the result still renders as a bare object.
553        let rows = vec![json!({"id": 1, "name": "a"})];
554        let (out, singular) =
555            unwrap_rpc_rows(rows.clone(), &composite_call_plan("get_user", false));
556        assert!(singular);
557        assert_eq!(out, rows);
558    }
559
560    #[test]
561    fn leaves_single_key_row_untouched_when_key_is_not_function_name() {
562        // A single real column that happens to be the only column should not be
563        // mistaken for the function-name wrapper.
564        let rows = vec![json!({"id": 7})];
565        let (out, _) = unwrap_rpc_rows(rows.clone(), &call_plan("get_thing", false));
566        assert_eq!(out, rows);
567    }
568}
postrust_server/app.rs

postrust_server/
app.rs
