feat(typedsql): support column & param nullability (prisma#4979)

libs/test-cli/src/main.rs

@@ -212,7 +212,6 @@ async fn main() -> anyhow::Result<()> {
  .first_datasource()
  .load_url(|key| std::env::var(key).ok())
  .unwrap(),
- force: false,
  queries: vec![SqlQueryInput {
  name: "query".to_string(),
  source: query_str,

quaint/src/connector.rs

@@ -12,11 +12,11 @@
 mod column_type;
 mod connection_info;
 
+mod describe;
 pub mod external;
 pub mod metrics;
 #[cfg(native)]
 pub mod native;
-mod parsed_query;
 mod queryable;
 mod result_set;
 #[cfg(any(feature = "mssql-native", feature = "postgresql-native", feature = "mysql-native"))]
@@ -32,8 +32,8 @@ pub use connection_info::*;
 #[cfg(native)]
 pub use native::*;
 
+pub use describe::*;
 pub use external::*;
-pub use parsed_query::*;
 pub use queryable::*;
 pub use transaction::*;
 

quaint/src/connector/parsed_query.rs → quaint/src/connector/describe.rs

@@ -3,26 +3,34 @@ use std::borrow::Cow;
 use super::ColumnType;
 
 #[derive(Debug)]
-pub struct ParsedRawQuery {
- pub parameters: Vec<ParsedRawParameter>,
- pub columns: Vec<ParsedRawColumn>,
+pub struct DescribedQuery {
+ pub parameters: Vec<DescribedParameter>,
+ pub columns: Vec<DescribedColumn>,
+ pub enum_names: Option<Vec<String>>,
+}
+
+impl DescribedQuery {
+ pub fn param_enum_names(&self) -> Vec<&str> {
+ self.parameters.iter().filter_map(|p| p.enum_name.as_deref()).collect()
+ }
 }
 
 #[derive(Debug)]
-pub struct ParsedRawParameter {
+pub struct DescribedParameter {
  pub name: String,
  pub typ: ColumnType,
  pub enum_name: Option<String>,
 }
 
 #[derive(Debug)]
-pub struct ParsedRawColumn {
+pub struct DescribedColumn {
  pub name: String,
  pub typ: ColumnType,
+ pub nullable: bool,
  pub enum_name: Option<String>,
 }
 
-impl ParsedRawParameter {
+impl DescribedParameter {
  pub fn new_named<'a>(name: impl Into<Cow<'a, str>>, typ: impl Into<ColumnType>) -> Self {
  let name: Cow<'_, str> = name.into();
 
@@ -52,14 +60,15 @@ impl ParsedRawParameter {
  }
 }
 
-impl ParsedRawColumn {
+impl DescribedColumn {
  pub fn new_named<'a>(name: impl Into<Cow<'a, str>>, typ: impl Into<ColumnType>) -> Self {
  let name: Cow<'_, str> = name.into();
 
  Self {
  name: name.into_owned(),
  typ: typ.into(),
  enum_name: None,
+ nullable: false,
  }
  }
 
@@ -68,11 +77,17 @@ impl ParsedRawColumn {
  name: format!("_{idx}"),
  typ: typ.into(),
  enum_name: None,
+ nullable: false,
  }
  }
 
  pub fn with_enum_name(mut self, enum_name: Option<String>) -> Self {
  self.enum_name = enum_name;
  self
  }
+
+ pub fn is_nullable(mut self, nullable: bool) -> Self {
+ self.nullable = nullable;
+ self
+ }
 }

quaint/src/connector/mssql/native/mod.rs

@@ -6,7 +6,7 @@ mod conversion;
 mod error;
 
 pub(crate) use crate::connector::mssql::MssqlUrl;
-use crate::connector::{timeout, IsolationLevel, ParsedRawQuery, Transaction, TransactionOptions};
+use crate::connector::{timeout, DescribedQuery, IsolationLevel, Transaction, TransactionOptions};
 
 use crate::{
  ast::{Query, Value},
@@ -183,8 +183,8 @@ impl Queryable for Mssql {
  self.query_raw(sql, params).await
  }
 
- async fn parse_raw_query(&self, _sql: &str) -> crate::Result<ParsedRawQuery> {
- unimplemented!("SQL Server support for raw query parsing is not implemented yet.")
+ async fn describe_query(&self, _sql: &str) -> crate::Result<DescribedQuery> {
+ unimplemented!("SQL Server does not support describe_query yet.")
  }
 
  async fn execute(&self, q: Query<'_>) -> crate::Result<u64> {

quaint/src/connector/mysql/native/mod.rs

@@ -6,7 +6,7 @@ mod conversion;
 mod error;
 
 pub(crate) use crate::connector::mysql::MysqlUrl;
-use crate::connector::{timeout, ColumnType, IsolationLevel, ParsedRawColumn, ParsedRawParameter, ParsedRawQuery};
+use crate::connector::{timeout, ColumnType, DescribedColumn, DescribedParameter, DescribedQuery, IsolationLevel};
 
 use crate::{
  ast::{Query, Value},
@@ -16,6 +16,7 @@ use crate::{
 };
 use async_trait::async_trait;
 use lru_cache::LruCache;
+use mysql_async::consts::ColumnFlags;
 use mysql_async::{
  self as my,
  prelude::{Query as _, Queryable as _},
@@ -247,21 +248,28 @@ impl Queryable for Mysql {
  self.query_raw(sql, params).await
  }
 
- async fn parse_raw_query(&self, sql: &str) -> crate::Result<ParsedRawQuery> {
+ async fn describe_query(&self, sql: &str) -> crate::Result<DescribedQuery> {
  self.prepared(sql, |stmt| async move {
  let columns = stmt
  .columns()
  .iter()
- .map(|col| ParsedRawColumn::new_named(col.name_str(), col))
+ .map(|col| {
+ DescribedColumn::new_named(col.name_str(), col)
+ .is_nullable(!col.flags().contains(ColumnFlags::NOT_NULL_FLAG))
+ })
  .collect();
  let parameters = stmt
  .params()
  .iter()
  .enumerate()
- .map(|(idx, col)| ParsedRawParameter::new_unnamed(idx, col))
+ .map(|(idx, col)| DescribedParameter::new_unnamed(idx, col))
  .collect();
 
- Ok(ParsedRawQuery { columns, parameters })
+ Ok(DescribedQuery {
+ columns,
+ parameters,
+ enum_names: None,
+ })
  })
  .await
  }

quaint/src/connector/postgres/native/explain.rs

@@ -0,0 +1,57 @@
+#[derive(serde::Deserialize, Debug)]
+#[serde(untagged)]
+pub(crate) enum Explain {
+ // NOTE: the returned JSON may not contain a `plan` field, for example, with `CALL` statements:
+ // https://github.com/launchbadge/sqlx/issues/1449
+ //
+ // In this case, we should just fall back to assuming all is nullable.
+ //
+ // It may also contain additional fields we don't care about, which should not break parsing:
+ // https://github.com/launchbadge/sqlx/issues/2587
+ // https://github.com/launchbadge/sqlx/issues/2622
+ Plan {
+ #[serde(rename = "Plan")]
+ plan: Plan,
+ },
+
+ // This ensures that parsing never technically fails.
+ //
+ // We don't want to specifically expect `"Utility Statement"` because there might be other cases
+ // and we don't care unless it contains a query plan anyway.
+ Other(serde::de::IgnoredAny),
+}
+
+#[derive(serde::Deserialize, Debug)]
+pub(crate) struct Plan {
+ #[serde(rename = "Join Type")]
+ pub(crate) join_type: Option<String>,
+ #[serde(rename = "Parent Relationship")]
+ pub(crate) parent_relation: Option<String>,
+ #[serde(rename = "Output")]
+ pub(crate) output: Option<Vec<String>>,
+ #[serde(rename = "Plans")]
+ pub(crate) plans: Option<Vec<Plan>>,
+}
+
+pub(crate) fn visit_plan(plan: &Plan, outputs: &[String], nullables: &mut Vec<Option<bool>>) {
+ if let Some(plan_outputs) = &plan.output {
+ // all outputs of a Full Join must be marked nullable
+ // otherwise, all outputs of the inner half of an outer join must be marked nullable
+ if plan.join_type.as_deref() == Some("Full") || plan.parent_relation.as_deref() == Some("Inner") {
+ for output in plan_outputs {
+ if let Some(i) = outputs.iter().position(|o| o == output) {
+ // N.B. this may produce false positives but those don't cause runtime errors
+ nullables[i] = Some(true);
+ }
+ }
+ }
+ }
+
+ if let Some(plans) = &plan.plans {
+ if let Some("Left") | Some("Right") = plan.join_type.as_deref() {
+ for plan in plans {
+ visit_plan(plan, outputs, nullables);
+ }
+ }
+ }
+}