From 6fde5a2579dc2df369844eff46087ff811e1f518 Mon Sep 17 00:00:00 2001
From: Austin Bonander <austin.bonander@gmail.com>
Date: Fri, 27 Mar 2020 17:28:39 -0700
Subject: [PATCH] implement nullability on macro output, test postgres

---
 sqlx-core/src/error.rs                 |  7 +++++++
 sqlx-macros/src/query_macros/output.rs | 11 ++++++-----
 src/result_ext.rs                      |  6 ++++++
 tests/postgres-macros.rs               | 16 ++++++++--------
 4 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/sqlx-core/src/error.rs b/sqlx-core/src/error.rs
index cc51ccf6..f6ecf36a 100644
--- a/sqlx-core/src/error.rs
+++ b/sqlx-core/src/error.rs
@@ -185,6 +185,13 @@ impl From<ProtocolError<'_>> for Error {
     }
 }
 
+impl From<UnexpectedNullError> for Error {
+    #[inline]
+    fn from(err: UnexpectedNullError) -> Self {
+        Error::Decode(err.into())
+    }
+}
+
 #[cfg(feature = "tls")]
 #[cfg_attr(docsrs, doc(cfg(feature = "tls")))]
 impl From<async_native_tls::Error> for Error {
diff --git a/sqlx-macros/src/query_macros/output.rs b/sqlx-macros/src/query_macros/output.rs
index df7cb7b9..0d2bc270 100644
--- a/sqlx-macros/src/query_macros/output.rs
+++ b/sqlx-macros/src/query_macros/output.rs
@@ -44,7 +44,7 @@ pub fn columns_to_rust<DB: DatabaseExt>(describe: &Describe<DB>) -> crate::Resul
 
             let ident = parse_ident(name)?;
 
-            let type_ = if let Some(type_info) = &column.type_info {
+            let mut type_ = if let Some(type_info) = &column.type_info {
                 <DB as DatabaseExt>::return_type_for_id(&type_info).map_or_else(
                     || {
                         let message = if let Some(feature_gate) =
@@ -88,10 +88,11 @@ pub fn columns_to_rust<DB: DatabaseExt>(describe: &Describe<DB>) -> crate::Resul
                 .to_compile_error()
             };
 
-            Ok(RustColumn {
-                ident,
-                type_: type_,
-            })
+            if !column.non_null.unwrap_or(false) {
+                type_ = quote! { Option<#type_> };
+            }
+
+            Ok(RustColumn { ident, type_ })
         })
         .collect::<crate::Result<Vec<_>>>()
 }
diff --git a/src/result_ext.rs b/src/result_ext.rs
index 8de1c191..d19de073 100644
--- a/src/result_ext.rs
+++ b/src/result_ext.rs
@@ -10,6 +10,12 @@ impl<T> ResultExt<T> for crate::Result<T> {
     }
 }
 
+impl<T> ResultExt<T> for crate::Result<Option<T>> {
+    fn try_unwrap_optional(self) -> crate::Result<T> {
+        self?.ok_or_else(|| UnexpectedNullError.into())
+    }
+}
+
 impl<T> ResultExt<Option<T>> for crate::Result<T> {
     fn try_unwrap_optional(self) -> crate::Result<Option<T>> {
         match self {
diff --git a/tests/postgres-macros.rs b/tests/postgres-macros.rs
index c23efc21..07250cd1 100644
--- a/tests/postgres-macros.rs
+++ b/tests/postgres-macros.rs
@@ -35,12 +35,12 @@ async fn test_text_var_char_char_n() -> anyhow::Result<()> {
     let mut conn = connect().await?;
 
     // TEXT
-
+    // we cannot infer nullability from an expression
     let rec = sqlx::query!("SELECT 'Hello'::text as greeting")
         .fetch_one(&mut conn)
         .await?;
 
-    assert_eq!(rec.greeting, "Hello");
+    assert_eq!(rec.greeting.as_deref(), Some("Hello"));
 
     // VARCHAR(N)
 
@@ -48,7 +48,7 @@ async fn test_text_var_char_char_n() -> anyhow::Result<()> {
         .fetch_one(&mut conn)
         .await?;
 
-    assert_eq!(rec.greeting, "Hello");
+    assert_eq!(rec.greeting.as_deref(), Some("Hello"));
 
     // CHAR(N)
 
@@ -56,7 +56,7 @@ async fn test_text_var_char_char_n() -> anyhow::Result<()> {
         .fetch_one(&mut conn)
         .await?;
 
-    assert_eq!(rec.greeting, "Hello");
+    assert_eq!(rec.greeting.as_deref(), Some("Hello"));
 
     Ok(())
 }
@@ -164,7 +164,7 @@ async fn query_by_string() -> anyhow::Result<()> {
     .fetch_one(&mut conn)
     .await?;
 
-    assert_eq!(result.string, string);
+    assert_eq!(result.string, Some(string));
 
     Ok(())
 }
@@ -212,7 +212,7 @@ async fn test_many_args() -> anyhow::Result<()> {
         .await?;
 
     for (i, row) in rows.iter().enumerate() {
-        assert_eq!(i as i32, row.id);
+        assert_eq!(Some(i as i32), row.id);
     }
 
     Ok(())
@@ -229,7 +229,7 @@ async fn test_array_from_slice() -> anyhow::Result<()> {
         .fetch_one(&mut conn)
         .await?;
 
-    assert_eq!(result.my_array, vec![1, 2, 3, 4]);
+    assert_eq!(result.my_array, Some(vec![1, 2, 3, 4]));
 
     println!("result ID: {:?}", result.my_array);
 
@@ -237,7 +237,7 @@ async fn test_array_from_slice() -> anyhow::Result<()> {
         .fetch_one(&mut conn)
         .await?;
 
-    assert_eq!(account.my_array, vec![4, 3, 2, 1]);
+    assert_eq!(account.my_array, Some(vec![4, 3, 2, 1]));
 
     println!("account ID: {:?}", account.my_array);