From 23a0fd35b7751c4577113fdb4ce5917763c87fe9 Mon Sep 17 00:00:00 2001
From: Austin Bonander <austin@launchbadge.com>
Date: Mon, 31 Jul 2023 13:23:31 -0700
Subject: [PATCH] refactor(postgres): simplify query plan handling, add unit
 test

---
 sqlx-postgres/src/connection/describe.rs | 105 +++++++++++++++++++----
 1 file changed, 87 insertions(+), 18 deletions(-)

diff --git a/sqlx-postgres/src/connection/describe.rs b/sqlx-postgres/src/connection/describe.rs
index 53b47f8c8a..1cd28a8b54 100644
--- a/sqlx-postgres/src/connection/describe.rs
+++ b/sqlx-postgres/src/connection/describe.rs
@@ -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)
@@ -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>,
@@ -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:?}"
+    )
+}