From 1d82d6f50f1988a092c3b832398d9c830221644d Mon Sep 17 00:00:00 2001
From: mrl5 <31549762+mrl5@users.noreply.github.com>
Date: Mon, 17 Jul 2023 18:25:38 +0200
Subject: [PATCH] fix(postgres): sqlx prepare fails if
 shared_preload_libraries=pg_stat_statements

closes #2622

refs:
* https://serde.rs/enum-representations.html#untagged
* https://serde.rs/field-attrs.html#skip
* https://www.postgresql.org/docs/current/pgstatstatements.html
  * https://www.postgresql.org/docs/current/runtime-config-statistics.html#GUC-COMPUTE-QUERY-ID
---
 sqlx-postgres/src/connection/describe.rs | 28 +++++++++++++++---------
 tests/docker-compose.yml                 |  2 +-
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/sqlx-postgres/src/connection/describe.rs b/sqlx-postgres/src/connection/describe.rs
index bcfbf3a9cf..9375bce693 100644
--- a/sqlx-postgres/src/connection/describe.rs
+++ b/sqlx-postgres/src/connection/describe.rs
@@ -451,15 +451,11 @@ WHERE rngtypid = $1
 
         let mut nullables = Vec::new();
 
-        if let Explain::Plan(
-            plan @ Plan {
-                output: Some(outputs),
-                ..
-            },
-        ) = &explain
-        {
-            nullables.resize(outputs.len(), None);
-            visit_plan(&plan, outputs, &mut nullables);
+        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);
+            }
         }
 
         Ok(nullables)
@@ -492,15 +488,27 @@ fn visit_plan(plan: &Plan, outputs: &[String], nullables: &mut Vec<Option<bool>>
 }
 
 #[derive(serde::Deserialize)]
+#[serde(untagged)]
 enum Explain {
     /// {"Plan": ...} -- returned for most statements
-    Plan(Plan),
+    QueryPlan(QueryPlan),
     /// The string "Utility Statement" -- returned for
     /// a CALL statement
     #[serde(rename = "Utility Statement")]
     UtilityStatement,
 }
 
+#[derive(serde::Deserialize)]
+struct QueryPlan {
+    #[serde(rename = "Plan")]
+    plan: Plan,
+    /// present when either pg_stat_statements is loaded and/or compute_query_id is enabled
+    /// https://www.postgresql.org/docs/current/pgstatstatements.html
+    /// https://www.postgresql.org/docs/current/runtime-config-statistics.html#GUC-COMPUTE-QUERY-ID
+    #[serde(rename = "Query Identifier", skip)]
+    _query_identifier: Option<u64>,
+}
+
 #[derive(serde::Deserialize)]
 struct Plan {
     #[serde(rename = "Join Type")]
diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml
index f521efb3a4..d50e431a49 100644
--- a/tests/docker-compose.yml
+++ b/tests/docker-compose.yml
@@ -183,7 +183,7 @@ services:
         volumes:
             - "./postgres/setup.sql:/docker-entrypoint-initdb.d/setup.sql"
         command: >
-            -c ssl=on -c ssl_cert_file=/var/lib/postgresql/server.crt -c ssl_key_file=/var/lib/postgresql/server.key
+            -c ssl=on -c ssl_cert_file=/var/lib/postgresql/server.crt -c ssl_key_file=/var/lib/postgresql/server.key -c shared_preload_libraries=pg_stat_statements
 
     postgres_15_client_ssl:
         build: