Skip to content

Commit

Permalink
refactor(postgres): simplify query plan handling, add unit test
Browse files Browse the repository at this point in the history
  • Loading branch information
abonander committed Jul 31, 2023
1 parent 9102d5d commit 23a0fd3
Showing 1 changed file with 87 additions and 18 deletions.
105 changes: 87 additions & 18 deletions sqlx-postgres/src/connection/describe.rs
Expand Up @@ -451,11 +451,16 @@ WHERE rngtypid = $1

let mut nullables = Vec::new();

if let Explain::QueryPlan(query_plan @ QueryPlan { plan, .. }) = &explain {
if let Some(outputs) = &query_plan.plan.output {
nullables.resize(outputs.len(), None);
visit_plan(&plan, outputs, &mut nullables);
}
if let Explain::Plan {
plan:
plan @ Plan {
output: Some(ref outputs),
..
},
} = &explain
{
nullables.resize(outputs.len(), None);
visit_plan(plan, outputs, &mut nullables);
}

Ok(nullables)
Expand Down Expand Up @@ -487,23 +492,30 @@ fn visit_plan(plan: &Plan, outputs: &[String], nullables: &mut Vec<Option<bool>>
}
}

#[derive(serde::Deserialize)]
#[derive(serde::Deserialize, Debug)]
#[serde(untagged)]
enum Explain {
/// {"Plan": ...} -- returned for most statements
QueryPlan(QueryPlan),
/// The string "Utility Statement" -- returned for
/// a CALL statement
UtilityStatement(String),
}

#[derive(serde::Deserialize)]
struct QueryPlan {
#[serde(rename = "Plan")]
plan: Plan,
// NOTE: the returned JSON may not contain a `plan` field, for example, with `CALL` statements:
// https://github.com/launchbadge/sqlx/issues/1449
//
// In this case, we should just fall back to assuming all is nullable.
//
// It may also contain additional fields we don't care about, which should not break parsing:
// https://github.com/launchbadge/sqlx/issues/2587
// https://github.com/launchbadge/sqlx/issues/2622
Plan {
#[serde(rename = "Plan")]
plan: Plan,
},

// This ensures that parsing never technically fails.
//
// We don't want to specifically expect `"Utility Statement"` because there might be other cases
// and we don't care unless it contains a query plan anyway.
Other(serde::de::IgnoredAny),
}

#[derive(serde::Deserialize)]
#[derive(serde::Deserialize, Debug)]
struct Plan {
#[serde(rename = "Join Type")]
join_type: Option<String>,
Expand All @@ -514,3 +526,60 @@ struct Plan {
#[serde(rename = "Plans")]
plans: Option<Vec<Plan>>,
}

#[test]
fn explain_parsing() {
let normal_plan = r#"[
{
"Plan": {
"Node Type": "Result",
"Parallel Aware": false,
"Async Capable": false,
"Startup Cost": 0.00,
"Total Cost": 0.01,
"Plan Rows": 1,
"Plan Width": 4,
"Output": ["1"]
}
}
]"#;

// https://github.com/launchbadge/sqlx/issues/2622
let extra_field = r#"[
{
"Plan": {
"Node Type": "Result",
"Parallel Aware": false,
"Async Capable": false,
"Startup Cost": 0.00,
"Total Cost": 0.01,
"Plan Rows": 1,
"Plan Width": 4,
"Output": ["1"]
},
"Query Identifier": 1147616880456321454
}
]"#;

// https://github.com/launchbadge/sqlx/issues/1449
let utility_statement = r#"["Utility Statement"]"#;

let normal_plan_parsed = serde_json::from_str::<[Explain; 1]>(normal_plan).unwrap();
let extra_field_parsed = serde_json::from_str::<[Explain; 1]>(extra_field).unwrap();
let utility_statement_parsed = serde_json::from_str::<[Explain; 1]>(utility_statement).unwrap();

assert!(
matches!(normal_plan_parsed, [Explain::Plan { plan: Plan { .. } }]),
"unexpected parse from {normal_plan:?}: {normal_plan_parsed:?}"
);

assert!(
matches!(extra_field_parsed, [Explain::Plan { plan: Plan { .. } }]),
"unexpected parse from {extra_field:?}: {extra_field_parsed:?}"
);

assert!(
matches!(utility_statement_parsed, [Explain::Other(_)]),
"unexpected parse from {utility_statement:?}: {utility_statement_parsed:?}"
)
}

0 comments on commit 23a0fd3

Please sign in to comment.