feat(macros): type override annotations for columns

sqlx-macros/src/query/mod.rs

@@ -208,6 +208,15 @@ where
             RecordType::Generated => {
                 let record_name: Type = syn::parse_str("Record").unwrap();
 
+                for rust_col in &columns {
+                    if rust_col.type_.is_none() {
+                        return Err(
+                            "columns may not have wildcard overrides in `query!()` or `query_as!()"
+                                .into(),
+                        );
+                    }
+                }
+
                 let record_fields = columns.iter().map(
                     |&output::RustColumn {
                          ref ident,

sqlx-macros/src/query/output.rs

@@ -1,17 +1,19 @@
 use proc_macro2::{Ident, Span, TokenStream};
-use quote::quote;
+use quote::{quote, ToTokens};
 use syn::Type;
 
-use sqlx_core::describe::Describe;
+use sqlx_core::describe::{Column, Describe};
 
 use crate::database::DatabaseExt;
 
 use crate::query::QueryMacroInput;
 use std::fmt::{self, Display, Formatter};
+use syn::parse::{Parse, ParseStream};
+use syn::Token;
 
 pub struct RustColumn {
     pub(super) ident: Ident,
-    pub(super) type_: TokenStream,
+    pub(super) type_: Option<TokenStream>,
 }
 
 struct DisplayColumn<'a> {
@@ -20,6 +22,18 @@ struct DisplayColumn<'a> {
     name: &'a str,
 }
 
+struct ColumnDecl {
+    ident: Ident,
+    // TIL Rust still has OOP keywords like `abstract`, `final`, `override` and `virtual` reserved
+    r#override: Option<ColumnOverride>,
+}
+
+enum ColumnOverride {
+    NonNull,
+    Wildcard,
+    Exact(Type),
+}
+
 impl Display for DisplayColumn<'_> {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         write!(f, "column #{} ({:?})", self.idx + 1, self.name)
@@ -32,15 +46,81 @@ pub fn columns_to_rust<DB: DatabaseExt>(describe: &Describe<DB>) -> crate::Resul
         .iter()
         .enumerate()
         .map(|(i, column)| -> crate::Result<_> {
-            let name = &*column.name;
-            let ident = parse_ident(name)?;
+            // add raw prefix to all identifiers
+            let decl = ColumnDecl::parse(&column.name)
+                .map_err(|e| format!("column name {:?} is invalid: {}", column.name, e))?;
 
-            let mut type_ = if let Some(type_info) = &column.type_info {
+            let type_ = match decl.r#override {
+                Some(ColumnOverride::Exact(ty)) => Some(ty.to_token_stream()),
+                Some(ColumnOverride::Wildcard) => None,
+                Some(ColumnOverride::NonNull) => Some(get_column_type(i, column)),
+                None => {
+                    let type_ = get_column_type(i, column);
+
+                    if !column.not_null.unwrap_or(false) {
+                        Some(quote! { Option<#type_> })
+                    } else {
+                        Some(type_)
+                    }
+                }
+            };
+
+            Ok(RustColumn {
+                ident: decl.ident,
+                type_,
+            })
+        })
+        .collect::<crate::Result<Vec<_>>>()
+}
+
+pub fn quote_query_as<DB: DatabaseExt>(
+    input: &QueryMacroInput,
+    out_ty: &Type,
+    bind_args: &Ident,
+    columns: &[RustColumn],
+) -> TokenStream {
+    let instantiations = columns.iter().enumerate().map(
+        |(
+            i,
+            &RustColumn {
+                ref ident,
+                ref type_,
+                ..
+            },
+        )| {
+            match (input.checked, type_) {
+                // we guarantee the type is valid so we can skip the runtime check
+                (true, Some(type_)) => quote! {
+                    #ident: row.try_get_unchecked::<#type_, _>(#i).try_unwrap_optional()?
+                },
+                // type was overridden to be a wildcard so we fallback to the runtime check
+                (true, None) => quote! ( #ident: row.try_get(#i)? ),
+                // macro is the `_unchecked!()` variant so this will die in decoding if it's wrong
+                (false, _) => quote!( #ident: row.try_get_unchecked(#i)? ),
+            }
+        },
+    );
+
+    let db_path = DB::db_path();
+    let row_path = DB::row_path();
+    let sql = &input.src;
+
+    quote! {
+        sqlx::query_with::<#db_path, _>(#sql, #bind_args).try_map(|row: #row_path| {
+            use sqlx::Row as _;
+            use sqlx::result_ext::ResultExt as _;
+
+            Ok(#out_ty { #(#instantiations),* })
+        })
+    }
+}
+
+fn get_column_type<DB: DatabaseExt>(i: usize, column: &Column<DB>) -> TokenStream {
+    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) =
-                            <DB as DatabaseExt>::get_feature_gate(&type_info)
-                        {
+                let message =
+                    if let Some(feature_gate) = <DB as DatabaseExt>::get_feature_gate(&type_info) {
                         format!(
                             "optional feature `{feat}` required for type {ty} of {col}",
                             ty = &type_info,
@@ -77,57 +157,57 @@ pub fn columns_to_rust<DB: DatabaseExt>(describe: &Describe<DB>) -> crate::Resul
             ),
         )
         .to_compile_error()
-            };
-
-            if !column.not_null.unwrap_or(false) {
-                type_ = quote! { Option<#type_> };
     }
-
-            Ok(RustColumn { ident, type_ })
-        })
-        .collect::<crate::Result<Vec<_>>>()
 }
 
-pub fn quote_query_as<DB: DatabaseExt>(
-    input: &QueryMacroInput,
-    out_ty: &Type,
-    bind_args: &Ident,
-    columns: &[RustColumn],
-) -> TokenStream {
-    let instantiations = columns.iter().enumerate().map(
-        |(
-            i,
-            &RustColumn {
-                ref ident,
-                ref type_,
-                ..
-            },
-        )| {
-            // For "checked" queries, the macro checks these at compile time and using "try_get"
-            // would also perform pointless runtime checks
+impl ColumnDecl {
+    fn parse(col_name: &str) -> crate::Result<Self> {
+        // find the end of the identifier because we want to use our own logic to parse it
+        // if we tried to feed this into `syn::parse_str()` we might get an un-great error
+        // for some kinds of invalid identifiers
+        let (ident, remainder) = if let Some(i) = col_name.find(&[':', '!'][..]) {
+            let (ident, remainder) = col_name.split_at(i);
 
-            if input.checked {
-                quote!( #ident: row.try_get_unchecked::<#type_, _>(#i).try_unwrap_optional()? )
+            (parse_ident(ident)?, remainder)
         } else {
-                quote!( #ident: row.try_get_unchecked(#i)? )
-            }
+            (parse_ident(col_name)?, "")
+        };
+
+        Ok(ColumnDecl {
+            ident,
+            r#override: if !remainder.is_empty() {
+                Some(syn::parse_str(remainder)?)
+            } else {
+                None
             },
-    );
-
-    let db_path = DB::db_path();
-    let row_path = DB::row_path();
-    let sql = &input.src;
-
-    quote! {
-        sqlx::query_with::<#db_path, _>(#sql, #bind_args).try_map(|row: #row_path| {
-            use sqlx::Row as _;
-            use sqlx::result_ext::ResultExt as _;
-
-            Ok(#out_ty { #(#instantiations),* })
         })
     }
 }
 
+impl Parse for ColumnOverride {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let lookahead = input.lookahead1();
+
+        if lookahead.peek(Token![:]) {
+            input.parse::<Token![:]>()?;
+
+            let ty = Type::parse(input)?;
+
+            if let Type::Infer(_) = ty {
+                Ok(ColumnOverride::Wildcard)
+            } else {
+                Ok(ColumnOverride::Exact(ty))
+            }
+        } else if lookahead.peek(Token![!]) {
+            input.parse::<Token![!]>()?;
+
+            Ok(ColumnOverride::NonNull)
+        } else {
+            Err(lookahead.error())
+        }
+    }
+}
+
 fn parse_ident(name: &str) -> crate::Result<Ident> {
     // workaround for the following issue (it's semi-fixed but still spits out extra diagnostics)
     // https://github.com/dtolnay/syn/issues/749#issuecomment-575451318

tests/postgres/macros.rs

@@ -254,3 +254,50 @@ async fn fetch_is_usable_issue_224() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[sqlx_macros::test]
+async fn test_column_override_not_null() -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+
+    let record = sqlx::query!(r#"select 1 as "id!""#)
+        .fetch_one(&mut conn)
+        .await?;
+
+    assert_eq!(record.id, 1);
+
+    Ok(())
+}
+
+#[derive(PartialEq, Eq, Debug, sqlx::Type)]
+#[sqlx(transparent)]
+struct MyInt4(i32);
+
+#[sqlx_macros::test]
+async fn test_column_override_wildcard() -> anyhow::Result<()> {
+    struct Record {
+        id: MyInt4,
+    }
+
+    let mut conn = new::<Postgres>().await?;
+
+    let record = sqlx::query_as!(Record, r#"select 1 as "id: _""#)
+        .fetch_one(&mut conn)
+        .await?;
+
+    assert_eq!(record.id, MyInt4(1));
+
+    Ok(())
+}
+
+#[sqlx_macros::test]
+async fn test_column_override_exact() -> anyhow::Result<()> {
+    let mut conn = new::<Postgres>().await?;
+
+    let record = sqlx::query!(r#"select 1 as "id: MyInt4""#)
+        .fetch_one(&mut conn)
+        .await?;
+
+    assert_eq!(record.id, MyInt4(1));
+
+    Ok(())
+}