diff --git a/src/db.rs b/src/db.rs index f2cefb9..f983882 100644 --- a/src/db.rs +++ b/src/db.rs @@ -3,6 +3,7 @@ use anyhow::{Context, Result}; use native_tls::TlsConnector; use postgres_native_tls::MakeTlsConnector; use serde_json::{Map, Value}; +use thiserror::Error; use tokio_postgres::{Client, Row}; #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] @@ -95,13 +96,35 @@ pub struct QueryExecution { pub truncated: bool, } +#[derive(Debug, Error)] +pub enum QueryExecutionError { + #[error("{source}")] + Connection { source: anyhow::Error }, + #[error("{context}")] + Session { + context: &'static str, + source: tokio_postgres::Error, + details: Vec, + }, + #[error("{message}")] + Query { + message: String, + hint: Option, + details: Vec, + }, +} + pub async fn execute_readonly_query( profile: &DatabaseProfile, sql: &str, has_limit: bool, -) -> Result { - let url = profile.connection_url()?; - let mut client = connect(&url).await?; +) -> std::result::Result { + let url = profile + .connection_url() + .map_err(|source| QueryExecutionError::Connection { source })?; + let mut client = connect(&url) + .await + .map_err(|source| QueryExecutionError::Connection { source })?; run_in_readonly_session(&mut client, profile, sql, has_limit).await } @@ -111,13 +134,15 @@ async fn run_in_readonly_session( profile: &DatabaseProfile, sql: &str, has_limit: bool, -) -> Result { +) -> std::result::Result { let transaction = client .build_transaction() .read_only(true) .start() .await - .context("failed to start read-only transaction")?; + .map_err(|source| { + session_execution_error("failed to start read-only transaction", source) + })?; transaction .execute( @@ -128,7 +153,7 @@ async fn run_in_readonly_session( &[], ) .await - .context("failed to set statement_timeout")?; + .map_err(|source| session_execution_error("failed to set statement_timeout", source))?; let effective_sql = if has_limit { sql.to_string() @@ -142,12 +167,12 @@ async fn run_in_readonly_session( let rows = transaction .query(&effective_sql, &[]) .await - .context("query execution failed")?; + .map_err(query_execution_error)?; transaction .rollback() .await - .context("failed to rollback transaction")?; + .map_err(|source| session_execution_error("failed to rollback transaction", source))?; let mut truncated = false; let mut result_rows = rows; @@ -178,6 +203,45 @@ async fn run_in_readonly_session( }) } +fn query_execution_error(source: tokio_postgres::Error) -> QueryExecutionError { + match source.as_db_error() { + Some(db_error) => QueryExecutionError::Query { + message: db_error.message().to_string(), + hint: db_error.hint().map(ToString::to_string), + details: database_error_details(db_error), + }, + None => QueryExecutionError::Query { + message: source.to_string(), + hint: None, + details: Vec::new(), + }, + } +} + +fn session_execution_error( + context: &'static str, + source: tokio_postgres::Error, +) -> QueryExecutionError { + let details = source + .as_db_error() + .map(database_error_details) + .unwrap_or_default(); + + QueryExecutionError::Session { + context, + source, + details, + } +} + +fn database_error_details(db_error: &tokio_postgres::error::DbError) -> Vec { + let mut details = vec![format!("SQLSTATE {}", db_error.code().code())]; + if let Some(detail) = db_error.detail() { + details.push(detail.to_string()); + } + details +} + fn row_to_json_map(row: &Row) -> Map { let mut map = Map::new(); for (i, column) in row.columns().iter().enumerate() { diff --git a/src/main.rs b/src/main.rs index 0ef3f56..f6ef633 100644 --- a/src/main.rs +++ b/src/main.rs @@ -385,19 +385,7 @@ async fn cmd_run( { Ok(r) => r, Err(e) => { - print_error( - ErrorBody { - code: "database_error".to_string(), - message: e.to_string(), - hint: Some(format!( - "Check that {} is set and the database is reachable", - profile.url_env - )), - location: None, - details: vec![], - }, - pretty, - ); + print_error(database_error_body(e, &profile.url_env), pretty); } }; @@ -414,6 +402,47 @@ async fn cmd_run( Ok(()) } +fn database_error_body( + error: querygate::db::QueryExecutionError, + url_env: &str, +) -> ErrorBody { + let reachability_hint = || { + Some(format!( + "Check that {url_env} is set and the database is reachable" + )) + }; + + match error { + querygate::db::QueryExecutionError::Query { + message, + hint, + details, + } => ErrorBody { + code: "database_error".to_string(), + message, + hint, + location: None, + details, + }, + querygate::db::QueryExecutionError::Connection { source } => ErrorBody { + code: "database_error".to_string(), + message: source.to_string(), + hint: reachability_hint(), + location: None, + details: vec![], + }, + querygate::db::QueryExecutionError::Session { + context, details, .. + } => ErrorBody { + code: "database_error".to_string(), + message: context.to_string(), + hint: reachability_hint(), + location: None, + details, + }, + } +} + struct SqlValidation { has_limit: bool, } @@ -515,3 +544,75 @@ fn read_sql(sql_arg: Option) -> Result { } Ok(trimmed.to_string()) } + +#[cfg(test)] +mod tests { + use super::database_error_body; + + #[test] + fn query_errors_surface_database_payload_without_reachability_hint() { + let body = database_error_body( + querygate::db::QueryExecutionError::Query { + message: "column \"missing\" does not exist".to_string(), + hint: Some( + "Perhaps you meant to reference the column \"users.name\".".to_string(), + ), + details: vec![ + "SQLSTATE 42703".to_string(), + "There is no column named missing.".to_string(), + ], + }, + "QUERYGATE_DATABASE_URL", + ); + + assert_eq!(body.code, "database_error"); + assert_eq!(body.message, "column \"missing\" does not exist"); + assert_eq!( + body.hint, + Some("Perhaps you meant to reference the column \"users.name\".".to_string()) + ); + assert_eq!( + body.details, + vec![ + "SQLSTATE 42703".to_string(), + "There is no column named missing.".to_string(), + ] + ); + } + + #[test] + fn query_errors_without_database_hint_do_not_use_reachability_hint() { + let body = database_error_body( + querygate::db::QueryExecutionError::Query { + message: "relation \"missing\" does not exist".to_string(), + hint: None, + details: vec!["SQLSTATE 42P01".to_string()], + }, + "QUERYGATE_DATABASE_URL", + ); + + assert_eq!(body.message, "relation \"missing\" does not exist"); + assert_eq!(body.hint, None); + assert_eq!(body.details, vec!["SQLSTATE 42P01".to_string()]); + } + + #[test] + fn connection_errors_keep_reachability_hint() { + let body = database_error_body( + querygate::db::QueryExecutionError::Connection { + source: anyhow::anyhow!("failed to connect to database"), + }, + "QUERYGATE_DATABASE_URL", + ); + + assert_eq!(body.message, "failed to connect to database"); + assert_eq!( + body.hint, + Some( + "Check that QUERYGATE_DATABASE_URL is set and the database is reachable" + .to_string() + ) + ); + assert!(body.details.is_empty()); + } +}