WIP support nullable columns in query!()

sqlx-core/src/describe.rs

@@ -39,6 +39,8 @@ where
     pub name: Option<Box<str>>,
     pub table_id: Option<DB::TableId>,
     pub type_info: DB::TypeInfo,
+    /// Whether or not the column may be `NULL` (or if that is even known).
+    pub nullability: Nullability,
 }
 
 impl<DB> Debug for Column<DB>
@@ -55,3 +57,10 @@ where
             .finish()
     }
 }
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum Nullability {
+    NonNull,
+    Nullable,
+    Unknown
+}

sqlx-core/src/postgres/executor.rs

@@ -1,13 +1,19 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
+use std::fmt::Write;
 use std::io;
 use std::sync::Arc;
 
 use futures_core::future::BoxFuture;
 use futures_core::stream::BoxStream;
+use futures_util::{stream, FutureExt, StreamExt, TryStreamExt};
 
-use crate::describe::{Column, Describe};
-use crate::postgres::protocol::{self, Encode, Message, StatementId, TypeFormat};
+use crate::arguments::Arguments;
+use crate::describe::{Column, Describe, Nullability};
+use crate::encode::IsNull::No;
 use crate::postgres::{PgArguments, PgRow, PgTypeInfo, Postgres};
+use crate::postgres::protocol::{self, Encode, Field, Message, StatementId, TypeFormat, TypeId};
+use crate::row::Row;
+use crate::postgres::types::SharedStr;
 
 #[derive(Debug)]
 enum Step {
@@ -31,7 +37,7 @@ impl super::PgConnection {
                 query,
                 param_types: &*args.types,
             }
-            .encode(self.stream.buffer_mut());
+                .encode(self.stream.buffer_mut());
 
             self.statement_cache.put(query.to_owned(), id);
 
@@ -53,7 +59,7 @@ impl super::PgConnection {
             values: &*args.values,
             result_formats: &[TypeFormat::Binary],
         }
-        .encode(self.stream.buffer_mut());
+            .encode(self.stream.buffer_mut());
     }
 
     fn write_execute(&mut self, portal: &str, limit: i32) {
@@ -280,27 +286,114 @@ impl super::PgConnection {
             }
         };
 
+        while let Some(_) = self.step().await? {}
+
+        let result_fields = result.map_or_else(Default::default, |r| r.fields);
+
+        // TODO: cache this result
+        let type_names = self.get_type_names(
+            params
+                .ids
+                .iter()
+                .cloned()
+                .chain(result_fields.iter().map(|field| field.type_id))
+        )
+            .await?;
+
         Ok(Describe {
             param_types: params
                 .ids
                 .iter()
-                .map(|id| PgTypeInfo::new(*id))
+                .map(|id| PgTypeInfo::new(*id, &type_names[&id.0]))
                 .collect::<Vec<_>>()
                 .into_boxed_slice(),
-            result_columns: result
-                .map(|r| r.fields)
-                .unwrap_or_default()
-                .into_vec()
-                .into_iter()
-                // TODO: Should [Column] just wrap [protocol::Field] ?
-                .map(|field| Column {
+            result_columns: self.map_result_columns(result_fields, type_names).await?
+                .into_boxed_slice(),
+        })
+    }
+
+    async fn get_type_names(&mut self, ids: impl IntoIterator<Item = TypeId>) -> crate::Result<HashMap<u32, SharedStr>> {
+        let type_ids: HashSet<u32> = ids.into_iter().map(|id| id.0).collect::<HashSet<u32>>();
+
+        let mut query = "select types.type_id, pg_type.typname from (VALUES ".to_string();
+        let mut args = PgArguments::default();
+        let mut pushed = false;
+
+        // TODO: dedup this with the one below, ideally as an API we can export
+        for (i, (&type_id, bind)) in type_ids.iter().zip((1 .. ).step_by(2)).enumerate() {
+            if pushed {
+                query += ", ";
+            }
+
+            pushed = true;
+            let _ = write!(query, "(${}, ${})", bind, bind + 1);
+
+            // not used in the output but ensures are values are sorted correctly
+            args.add(i as i32);
+            args.add(type_id as i32);
+        }
+
+        query += ") as types(idx, type_id) \
+        inner join pg_catalog.pg_type on pg_type.oid = type_id \
+        order by types.idx";
+
+        self.fetch(&query, args)
+            .map_ok(|row: PgRow| -> (u32, SharedStr) {
+                (row.get::<i32, _>(0) as u32, row.get::<String, _>(1).into())
+            })
+            .try_collect()
+            .await
+    }
+
+    async fn map_result_columns(&mut self, fields: Box<[Field]>, type_names: HashMap<u32, SharedStr>) -> crate::Result<Vec<Column<Postgres>>> {
+        use crate::describe::Nullability::*;
+
+        if fields.is_empty() { return Ok(vec![]); }
+
+        let mut query = "select col.idx, pg_attribute.attnotnull from (VALUES ".to_string();
+        let mut pushed = false;
+        let mut args = PgArguments::default();
+
+        for (i, (field, bind)) in fields.iter().zip((1 ..).step_by(3)).enumerate() {
+            if pushed {
+                query += ", ";
+            }
+
+            pushed = true;
+            let _ = write!(query, "(${}, ${}, ${})", bind, bind + 1, bind + 2);
+
+            args.add(i as i32);
+            args.add(field.table_id.map(|id| id as i32));
+            args.add(field.column_id);
+        }
+
+        query += ") as col(idx, table_id, col_idx) \
+        left join pg_catalog.pg_attribute on table_id is not null and attrelid = table_id and attnum = col_idx \
+        order by col.idx;";
+
+        log::trace!("describe pg_attribute query: {:#?}", query);
+
+        self.fetch(&query, args)
+            .zip(stream::iter(fields.into_vec().into_iter().enumerate()))
+            .map(|(row, (fidx, field))| -> crate::Result<Column<_>> {
+                let row = row?;
+                let idx = row.get::<i32, _>(0);
+                let nonnull = row.get::<Option<bool>, _>(1);
+
+                if idx != fidx as i32 {
+                    return Err(protocol_err!("missing field from query, field: {:?}", field).into());
+                }
+
+                Ok(Column {
                     name: field.name,
                     table_id: field.table_id,
-                    type_info: PgTypeInfo::new(field.type_id),
+                    type_info: PgTypeInfo::new(field.type_id, &type_names[&field.type_id.0]),
+                    nullability: nonnull.map(|nonnull| if nonnull { NonNull } else { Nullable })
+                        .unwrap_or(Unknown),
                 })
-                .collect::<Vec<_>>()
-                .into_boxed_slice(),
-        })
+            })
+            .try_collect()
+            .await
     }
 }
 

sqlx-core/src/postgres/types/bool.rs

@@ -7,13 +7,13 @@ use crate::types::HasSqlType;
 
 impl HasSqlType<bool> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::BOOL)
+        PgTypeInfo::new(TypeId::BOOL, "bool")
     }
 }
 
 impl HasSqlType<[bool]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_BOOL)
+        PgTypeInfo::new(TypeId::ARRAY_BOOL, "bool[]")
     }
 }
 

sqlx-core/src/postgres/types/bytes.rs

@@ -7,13 +7,13 @@ use crate::types::HasSqlType;
 
 impl HasSqlType<[u8]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::BYTEA)
+        PgTypeInfo::new(TypeId::BYTEA, "bytea")
     }
 }
 
 impl HasSqlType<[&'_ [u8]]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_BYTEA)
+        PgTypeInfo::new(TypeId::ARRAY_BYTEA, "bytea[]")
     }
 }
 

sqlx-core/src/postgres/types/chrono.rs

@@ -12,19 +12,19 @@ use crate::types::HasSqlType;
 
 impl HasSqlType<NaiveTime> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::TIME)
+        PgTypeInfo::new(TypeId::TIME, "time")
     }
 }
 
 impl HasSqlType<NaiveDate> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::DATE)
+        PgTypeInfo::new(TypeId::DATE, "date")
     }
 }
 
 impl HasSqlType<NaiveDateTime> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::TIMESTAMP)
+        PgTypeInfo::new(TypeId::TIMESTAMP, "timestamp")
     }
 }
 
@@ -33,25 +33,25 @@ where
     Tz: TimeZone,
 {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::TIMESTAMPTZ)
+        PgTypeInfo::new(TypeId::TIMESTAMPTZ, "timestamptz")
     }
 }
 
 impl HasSqlType<[NaiveTime]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_TIME)
+        PgTypeInfo::new(TypeId::ARRAY_TIME, "time[]")
     }
 }
 
 impl HasSqlType<[NaiveDate]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_DATE)
+        PgTypeInfo::new(TypeId::ARRAY_DATE, "date[]")
     }
 }
 
 impl HasSqlType<[NaiveDateTime]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_TIMESTAMP)
+        PgTypeInfo::new(TypeId::ARRAY_TIMESTAMP, "timestamp[]")
     }
 }
 
@@ -60,7 +60,7 @@ where
     Tz: TimeZone,
 {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_TIMESTAMPTZ)
+        PgTypeInfo::new(TypeId::ARRAY_TIMESTAMPTZ, "timestamp[]")
     }
 }
 

sqlx-core/src/postgres/types/float.rs

@@ -7,13 +7,13 @@ use crate::types::HasSqlType;
 
 impl HasSqlType<f32> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::FLOAT4)
+        PgTypeInfo::new(TypeId::FLOAT4, "float4")
     }
 }
 
 impl HasSqlType<[f32]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_FLOAT4)
+        PgTypeInfo::new(TypeId::ARRAY_FLOAT4, "float4[]")
     }
 }
 
@@ -33,13 +33,13 @@ impl Decode<Postgres> for f32 {
 
 impl HasSqlType<f64> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::FLOAT8)
+        PgTypeInfo::new(TypeId::FLOAT8, "float8")
     }
 }
 
 impl HasSqlType<[f64]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_FLOAT8)
+        PgTypeInfo::new(TypeId::ARRAY_FLOAT8, "float8[]")
     }
 }
 

sqlx-core/src/postgres/types/int.rs

@@ -9,13 +9,13 @@ use crate::types::HasSqlType;
 
 impl HasSqlType<i16> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::INT2)
+        PgTypeInfo::new(TypeId::INT2, "int2")
     }
 }
 
 impl HasSqlType<[i16]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_INT2)
+        PgTypeInfo::new(TypeId::ARRAY_INT2, "int2[]")
     }
 }
 
@@ -33,13 +33,13 @@ impl Decode<Postgres> for i16 {
 
 impl HasSqlType<i32> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::INT4)
+        PgTypeInfo::new(TypeId::INT4, "int4")
     }
 }
 
 impl HasSqlType<[i32]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_INT4)
+        PgTypeInfo::new(TypeId::ARRAY_INT4, "int4[]")
     }
 }
 
@@ -57,13 +57,13 @@ impl Decode<Postgres> for i32 {
 
 impl HasSqlType<i64> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::INT8)
+        PgTypeInfo::new(TypeId::INT8, "int8")
     }
 }
 
 impl HasSqlType<[i64]> for Postgres {
     fn type_info() -> PgTypeInfo {
-        PgTypeInfo::new(TypeId::ARRAY_INT8)
+        PgTypeInfo::new(TypeId::ARRAY_INT8, "int8[]")
     }
 }