From dd99e92513cfc53a46165648ff439b0097f2682b Mon Sep 17 00:00:00 2001
From: Flavian Desverne <flavian.dsv@gmail.com>
Date: Fri, 19 Jul 2024 16:51:14 +0200
Subject: [PATCH] fix regression test

---
 .../tests/new/regressions/prisma_6173.rs      |  4 +-
 .../sql-query-connector/src/ser_raw.rs        | 56 ++++++++++---------
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_6173.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_6173.rs
index 50b0d9f0d005..3aa1e9b0e2d2 100644
--- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_6173.rs
+++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_6173.rs
@@ -38,7 +38,9 @@ mod query_raw {
               "columns": [
                 "f0"
               ],
-              "types": [],
+              "types": [
+                "int"
+              ],
               "rows": [
                 [
                   "<rand_int>"
diff --git a/query-engine/connectors/sql-query-connector/src/ser_raw.rs b/query-engine/connectors/sql-query-connector/src/ser_raw.rs
index 048d5c3e5785..63b573737e2b 100644
--- a/query-engine/connectors/sql-query-connector/src/ser_raw.rs
+++ b/query-engine/connectors/sql-query-connector/src/ser_raw.rs
@@ -66,36 +66,42 @@ struct SerializedTypes<'a>(&'a ResultSet);
 
 impl<'a> SerializedTypes<'a> {
     fn infer_unknown_column_types(&self) -> Vec<ColumnType> {
-        let mut types = self.0.types().to_owned();
+        let rows = self.0;
+
+        if rows.is_empty() {
+            return Vec::with_capacity(0);
+        }
+
+        let row_len = rows.first().unwrap().len();
+
+        let mut types = vec![ColumnType::Unknown; row_len];
         let mut types_found = 0;
 
-        let unknown_indexes: Vec<_> = types
-            .iter()
-            .enumerate()
-            .filter_map(|(idx, ty)| match ty.is_unknown() {
-                true => Some(idx),
-                false => None,
-            })
-            .collect();
-
-        for idx in unknown_indexes.iter() {
-            // We attempt to infer types based on `quaint::Value` present in the rows if quaint couldn't infer a ColumnType.
-            // We need to go through every row because because empty and null arrays don't encode their inner type.
-            // In the best case scenario, this loop stops at the first row.
-            // In the worst case scenario, it'll keep looping until it finds an array with a non-null value.
-            for row in self.0.iter() {
-                let current_type = types[*idx];
-                let inferred_type = ColumnType::from(&row.at(*idx).unwrap().typed);
-
-                if current_type.is_unknown() && !inferred_type.is_unknown() {
-                    types[*idx] = inferred_type;
-                    types_found += 1;
-                    break;
+        'outer: for row in rows.iter() {
+            // While quaint already infers `ColumnType`s from the database, there are two cases where we need to rely on inferring from the actual result values:
+            // - When Quaint infers `ColumnType::Unknown`
+            // - When `query_raw` does not return columns in `ResultSet` when a call to a stored procedure is done (See https://github.com/prisma/prisma/issues/6173)
+            for (row_idx, value) in row.iter().enumerate() {
+                match rows.types().get(row_idx) {
+                    // If the ColumnType inferred from Quaint doesn't exist or it's unknown, infer the type from the value.
+                    Some(ColumnType::Unknown) | None => {
+                        let current_type = types[row_idx];
+                        let inferred_type = ColumnType::from(&value.typed);
+
+                        if current_type.is_unknown() && !inferred_type.is_unknown() {
+                            types[row_idx] = inferred_type;
+                            types_found += 1;
+                        }
+                    }
+                    Some(ty) => {
+                        types[row_idx] = *ty;
+                        types_found += 1;
+                    }
                 }
             }
 
-            if types_found == unknown_indexes.len() {
-                break;
+            if types_found == row_len {
+                break 'outer;
             }
         }