Warnings (#352)

edgedb-protocol/Cargo.toml

@@ -22,14 +22,15 @@ chrono = {version="0.4.23", optional=true, features=["std"], default-features=fa
 edgedb-errors = {path = "../edgedb-errors", version = "0.4.0" }
 bitflags = "2.4.0"
 serde = {version="1.0.190", features = ["derive"], optional=true}
+serde_json = {version="1", optional=true}
 
 [features]
 default = []
 with-num-bigint = ["num-bigint", "num-traits"]
 with-bigdecimal = ["bigdecimal", "num-bigint", "num-traits"]
 with-chrono = ["chrono"]
 all-types = ["with-num-bigint", "with-bigdecimal", "with-chrono"]
-with-serde = ["serde"]
+with-serde = ["serde", "serde_json"]
 
 [dev-dependencies]
 rand = "0.8"

edgedb-protocol/src/annotations.rs

@@ -0,0 +1,115 @@
+#[cfg(feature = "with-serde")]
+use crate::encoding::Annotations;
+
+/// CommandDataDescription1 may contain "warnings" annotations, whose value is
+/// a JSON array of this [Warning] type.
+#[derive(Debug, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "with-serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct Warning {
+    /// User-friendly explanation of the problem
+    pub message: String,
+
+    /// Name of the Python exception class
+    pub r#type: String,
+
+    /// Machine-friendly exception id
+    pub code: u64,
+
+    /// Name of the source file that caused the warning.
+    pub filename: Option<String>,
+
+    /// Additional user-friendly info
+    pub hint: Option<String>,
+
+    /// Developer-friendly explanation of why this problem occured
+    pub details: Option<String>,
+
+    /// Inclusive 0-based position within the source
+    #[cfg_attr(
+        feature = "with-serde",
+        serde(deserialize_with = "deserialize_usize_from_str")
+    )]
+    pub start: Option<usize>,
+
+    /// Exclusive 0-based position within the source
+    #[cfg_attr(
+        feature = "with-serde",
+        serde(deserialize_with = "deserialize_usize_from_str")
+    )]
+    pub end: Option<usize>,
+
+    /// 1-based index of the line of the start
+    #[cfg_attr(
+        feature = "with-serde",
+        serde(deserialize_with = "deserialize_usize_from_str")
+    )]
+    pub line: Option<usize>,
+
+    /// 1-based index of the column of the start
+    #[cfg_attr(
+        feature = "with-serde",
+        serde(deserialize_with = "deserialize_usize_from_str")
+    )]
+    pub col: Option<usize>,
+}
+
+impl std::fmt::Display for Warning {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let Warning {
+            filename,
+            line,
+            col,
+            r#type,
+            message,
+            ..
+        } = self;
+        let filename = filename
+            .as_ref()
+            .map(|f| format!("{f}:"))
+            .unwrap_or_default();
+        let line = line.clone().unwrap_or(1);
+        let col = col.clone().unwrap_or(1);
+
+        write!(f, "{type} at {filename}{line}:{col} {message}")
+    }
+}
+
+#[cfg(feature = "with-serde")]
+pub fn decode_warnings(annotations: &Annotations) -> Result<Vec<Warning>, edgedb_errors::Error> {
+    use edgedb_errors::{ErrorKind, ProtocolEncodingError};
+
+    const ANN_NAME: &str = "warnings";
+
+    if let Some(warnings) = annotations.get(ANN_NAME) {
+        serde_json::from_str::<Vec<_>>(warnings).map_err(|e| {
+            ProtocolEncodingError::with_source(e)
+                .context("Invalid JSON while decoding 'warnings' annotation")
+        })
+    } else {
+        Ok(vec![])
+    }
+}
+
+#[cfg(feature = "with-serde")]
+fn deserialize_usize_from_str<'de, D: serde::Deserializer<'de>>(
+    deserializer: D,
+) -> Result<Option<usize>, D::Error> {
+    use serde::Deserialize;
+
+    #[derive(Deserialize)]
+    #[serde(untagged)]
+    enum StringOrInt {
+        String(String),
+        Number(usize),
+        None,
+    }
+
+    match StringOrInt::deserialize(deserializer)? {
+        StringOrInt::String(s) => s
+            .parse::<usize>()
+            .map_err(serde::de::Error::custom)
+            .map(Some),
+        StringOrInt::Number(i) => Ok(Some(i)),
+        StringOrInt::None => Ok(None),
+    }
+}

edgedb-protocol/src/errors.rs

@@ -97,6 +97,11 @@ pub enum DecodeError {
     },
     #[snafu(display("missing required link or property"))]
     MissingRequiredElement { backtrace: Backtrace },
+    #[snafu(display("invalid format of {annotation} annotation"))]
+    InvalidAnnotationFormat {
+        backtrace: Backtrace,
+        annotation: &'static str,
+    },
 }
 
 #[derive(Snafu, Debug)]

edgedb-protocol/src/lib.rs

@@ -73,6 +73,7 @@ pub mod server_message;
 pub mod value;
 #[macro_use]
 pub mod value_opt;
+pub mod annotations;
 pub mod model;
 pub mod query_arg;
 

edgedb-tokio/Cargo.toml

@@ -11,7 +11,7 @@ readme = "README.md"
 rust-version.workspace = true
 
 [dependencies]
-edgedb-protocol = { path = "../edgedb-protocol", version = "0.6.0" }
+edgedb-protocol = { path = "../edgedb-protocol", version = "0.6.0", features = ["with-serde"] }
 edgedb-errors = { path = "../edgedb-errors", version = "0.4.1" }
 edgedb-derive = { path = "../edgedb-derive", version = "0.5.1", optional = true }
 tokio = { version = "1.15", features = ["net", "time", "sync", "macros"] }

edgedb-tokio/src/client.rs

@@ -11,11 +11,12 @@ use crate::builder::Config;
 use crate::errors::NoDataError;
 use crate::errors::{Error, ErrorKind, SHOULD_RETRY};
 use crate::options::{RetryOptions, TransactionOptions};
-use crate::raw::{Options, PoolState};
+use crate::raw::{Options, PoolState, Response};
 use crate::raw::{Pool, QueryCapabilities};
 use crate::state::{AliasesDelta, ConfigDelta, GlobalsDelta};
 use crate::state::{AliasesModifier, ConfigModifier, Fn, GlobalsModifier};
 use crate::transaction::{transaction, Transaction};
+use crate::ResultVerbose;
 
 /// The EdgeDB Client.
 ///
@@ -56,13 +57,14 @@ impl Client {
         Ok(())
     }
 
-    async fn query_and_retry<R, A>(
+    /// Query with retry.
+    async fn query_helper<R, A>(
         &self,
         query: impl AsRef<str>,
         arguments: &A,
         io_format: IoFormat,
         cardinality: Cardinality,
-    ) -> Result<Vec<R>, Error>
+    ) -> Result<Response<Vec<R>>, Error>
     where
         A: QueryArgs,
         R: QueryResult,
@@ -85,7 +87,7 @@ impl Client {
                 )
                 .await
             {
-                Ok(resp) => return Ok(resp.data),
+                Ok(resp) => return Ok(resp),
                 Err(e) => {
                     let allow_retry = match e.get::<QueryCapabilities>() {
                         // Error from a weird source, or just a bug
@@ -110,6 +112,31 @@ impl Client {
         }
     }
 
+    /// Execute a query and return a collection of results and warnings produced by the server.
+    ///
+    /// You will usually have to specify the return type for the query:
+    ///
+    /// ```rust,ignore
+    /// let greeting: (Vec<String>, _) = conn.query_with_warnings("select 'hello'", &()).await?;
+    /// ```
+    ///
+    /// This method can be used with both static arguments, like a tuple of
+    /// scalars, and with dynamic arguments [`edgedb_protocol::value::Value`].
+    /// Similarly, dynamically typed results are also supported.
+    pub async fn query_verbose<R, A>(
+        &self,
+        query: impl AsRef<str> + Send,
+        arguments: &A,
+    ) -> Result<ResultVerbose<Vec<R>>, Error>
+    where
+        A: QueryArgs,
+        R: QueryResult,
+    {
+        Client::query_helper(self, query, arguments, IoFormat::Binary, Cardinality::Many)
+            .await
+            .map(|Response { data, warnings, .. }| ResultVerbose { data, warnings })
+    }
+
     /// Execute a query and return a collection of results.
     ///
     /// You will usually have to specify the return type for the query:
@@ -134,7 +161,9 @@ impl Client {
         A: QueryArgs,
         R: QueryResult,
     {
-        Client::query_and_retry(self, query, arguments, IoFormat::Binary, Cardinality::Many).await
+        Client::query_helper(self, query, arguments, IoFormat::Binary, Cardinality::Many)
+            .await
+            .map(|r| r.data)
     }
 
     /// Execute a query and return a single result
@@ -162,15 +191,15 @@ impl Client {
         A: QueryArgs,
         R: QueryResult + Send,
     {
-        Client::query_and_retry(
+        Client::query_helper(
             self,
             query,
             arguments,
             IoFormat::Binary,
             Cardinality::AtMostOne,
         )
         .await
-        .map(|x| x.into_iter().next())
+        .map(|x| x.data.into_iter().next())
     }
 
     /// Execute a query and return a single result
@@ -207,7 +236,7 @@ impl Client {
         A: QueryArgs,
         R: QueryResult + Send,
     {
-        Client::query_and_retry(
+        Client::query_helper(
             self,
             query,
             arguments,
@@ -216,7 +245,8 @@ impl Client {
         )
         .await
         .and_then(|x| {
-            x.into_iter()
+            x.data
+                .into_iter()
                 .next()
                 .ok_or_else(|| NoDataError::with_message("query row returned zero results"))
         })
@@ -229,10 +259,11 @@ impl Client {
         arguments: &impl QueryArgs,
     ) -> Result<Json, Error> {
         let res = self
-            .query_and_retry::<String, _>(query, arguments, IoFormat::Json, Cardinality::Many)
+            .query_helper::<String, _>(query, arguments, IoFormat::Json, Cardinality::Many)
             .await?;
 
         let json = res
+            .data
             .into_iter()
             .next()
             .ok_or_else(|| NoDataError::with_message("query row returned zero results"))?;
@@ -266,11 +297,11 @@ impl Client {
         arguments: &impl QueryArgs,
     ) -> Result<Option<Json>, Error> {
         let res = self
-            .query_and_retry::<String, _>(query, arguments, IoFormat::Json, Cardinality::AtMostOne)
+            .query_helper::<String, _>(query, arguments, IoFormat::Json, Cardinality::AtMostOne)
             .await?;
 
         // we trust database to produce valid json
-        Ok(res.into_iter().next().map(Json::new_unchecked))
+        Ok(res.data.into_iter().next().map(Json::new_unchecked))
     }
 
     /// Execute a query and return a single result as JSON.

edgedb-tokio/src/lib.rs

@@ -146,7 +146,7 @@ pub use client::Client;
 pub use credentials::TlsSecurity;
 pub use errors::Error;
 pub use options::{RetryCondition, RetryOptions, TransactionOptions};
-pub use query_executor::QueryExecutor;
+pub use query_executor::{QueryExecutor, ResultVerbose};
 pub use state::{ConfigDelta, GlobalsDelta};
 pub use transaction::Transaction;
 

edgedb-tokio/src/query_executor.rs

@@ -1,10 +1,16 @@
-use edgedb_protocol::model::Json;
 use edgedb_protocol::query_arg::QueryArgs;
 use edgedb_protocol::QueryResult;
+use edgedb_protocol::{annotations::Warning, model::Json};
 use std::future::Future;
 
 use crate::{Client, Error, Transaction};
 
+#[non_exhaustive]
+pub struct ResultVerbose<R> {
+    pub data: R,
+    pub warnings: Vec<Warning>,
+}
+
 /// Abstracts over different query executors
 /// In particular &Client and &mut Transaction
 pub trait QueryExecutor: Sized {
@@ -18,6 +24,16 @@ pub trait QueryExecutor: Sized {
         A: QueryArgs,
         R: QueryResult + Send;
 
+    /// see [Client::query_with_warnings]
+    fn query_verbose<R, A>(
+        self,
+        query: impl AsRef<str> + Send,
+        arguments: &A,
+    ) -> impl Future<Output = Result<ResultVerbose<Vec<R>>, Error>> + Send
+    where
+        A: QueryArgs,
+        R: QueryResult + Send;
+
     /// see [Client::query_single]
     fn query_single<R, A>(
         self,
@@ -82,6 +98,18 @@ impl QueryExecutor for &Client {
         Client::query(self, query, arguments)
     }
 
+    fn query_verbose<R, A>(
+        self,
+        query: impl AsRef<str> + Send,
+        arguments: &A,
+    ) -> impl Future<Output = Result<ResultVerbose<Vec<R>>, Error>> + Send
+    where
+        A: QueryArgs,
+        R: QueryResult + Send,
+    {
+        Client::query_verbose(self, query, arguments)
+    }
+
     fn query_single<R, A>(
         self,
         query: impl AsRef<str> + Send,
@@ -151,6 +179,18 @@ impl QueryExecutor for &mut Transaction {
         Transaction::query(self, query, arguments)
     }
 
+    fn query_verbose<R, A>(
+        self,
+        query: impl AsRef<str> + Send,
+        arguments: &A,
+    ) -> impl Future<Output = Result<ResultVerbose<Vec<R>>, Error>> + Send
+    where
+        A: QueryArgs,
+        R: QueryResult + Send,
+    {
+        Transaction::query_verbose(self, query, arguments)
+    }
+
     fn query_single<R, A>(
         self,
         query: impl AsRef<str> + Send,