Accept Value::Str for Enum encoding (#308)

Co-authored-by: Aljaž Mur Eržen <[email protected]>
edgedb · Apr 9, 2024 · ca67c54 · ca67c54
1 parent 787a469
commit ca67c54
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 18 deletions.
diff --git a/edgedb-protocol/src/codec.rs b/edgedb-protocol/src/codec.rs
@@ -1364,11 +1364,12 @@ impl Codec for Enum {
         -> Result<(), EncodeError>
     {
         let val = match val {
-            Value::Enum(val) => val,
+            Value::Enum(val) => val.0.as_ref(),
+            Value::Str(val) => val.as_str(),
             _ => Err(errors::invalid_value(type_name::<Self>(), val))?,
         };
-        ensure!(self.members.get(&val.0).is_some(), errors::MissingEnumValue);
-        buf.extend(val.0.as_bytes());
+        ensure!(self.members.get(val).is_some(), errors::MissingEnumValue);
+        buf.extend(val.as_bytes());
         Ok(())
     }
 }
diff --git a/edgedb-protocol/src/query_arg.rs b/edgedb-protocol/src/query_arg.rs
@@ -224,21 +224,7 @@ impl QueryArg for Value {
             (Uuid(_), BaseScalar(d)) if d.id == codec::STD_UUID => Ok(()),
             (Enum(val), Enumeration(EnumerationTypeDescriptor { members, .. })) => {
                 let val = val.deref();
-                if members.iter().any(|c| c == val) {
-                    Ok(())
-                } else {
-                    let members = {
-                        let mut members = members
-                        .iter()
-                        .map(|c| format!("'{c}'"))
-                        .collect::<Vec<_>>();
-                        members.sort_unstable();
-                        members.join(", ")
-                    };
-                    Err(InvalidReferenceError::with_message(format!(
-                        "Expected one of: {members}, while enum value '{val}' was provided"
-                    )))
-                }
+                check_enum(val, members)
             }
             // TODO(tailhook) all types
             (_, desc) => Err(ctx.wrong_type(desc, self.kind())),
@@ -249,6 +235,22 @@ impl QueryArg for Value {
     }
 }
 
+pub(crate) fn check_enum(variant_name: &str, expected_members: &[String]) -> Result<(), Error> {
+    if expected_members.iter().any(|c| c == variant_name) {
+        Ok(())
+    } else {
+        let mut members = expected_members
+            .into_iter()
+            .map(|c| format!("'{c}'"))
+            .collect::<Vec<_>>();
+        members.sort_unstable();
+        let members = members.join(", ");
+        Err(InvalidReferenceError::with_message(format!(
+            "Expected one of: {members}, while enum value '{variant_name}' was provided"
+        )))
+    }
+}
+
 impl QueryArgs for Value {
     fn encode(&self, enc: &mut Encoder) -> Result<(), Error> {
         let codec = enc.ctx.build_codec()?;

diff --git a/edgedb-protocol/src/serialization/decode/raw_scalar.rs b/edgedb-protocol/src/serialization/decode/raw_scalar.rs
@@ -87,6 +87,11 @@ impl ScalarArg for &'_ str {
     fn check_descriptor(ctx: &DescriptorContext, pos: TypePos)
         -> Result<(), Error>
     {
+        // special case: &str can express an enum variant
+        if let Descriptor::Enumeration(_) = ctx.get(pos)? {
+            return Ok(())
+        }
+
         check_scalar(ctx, pos, String::uuid(), String::typename())
     }
     fn to_value(&self) -> Result<Value, Error> {

diff --git a/edgedb-tokio/tests/func/client.rs b/edgedb-tokio/tests/func/client.rs
@@ -1,3 +1,4 @@
+use edgedb_protocol::value::{EnumValue, Value};
 use edgedb_tokio::Client;
 use edgedb_errors::NoDataError;
 use futures_util::stream::{self, StreamExt};
@@ -42,6 +43,33 @@ async fn simple() -> anyhow::Result<()> {
     client.execute("SELECT 1+1", &()).await?;
     client.execute("START MIGRATION TO {}; ABORT MIGRATION", &()).await?;
 
+    // basic enum param
+    let enum_query = "SELECT <str>(<test::State>$0) = 'waiting'";
+    assert_eq!(
+        client.query_required_single::<bool, _>(
+            enum_query, &(Value::Enum(EnumValue::from("waiting")),)
+        ).await.unwrap(),
+        true
+    );
+
+    // unsupported: enum param as Value::Str
+    client.query_required_single::<bool, (Value, )>(
+        enum_query, &(Value::Str("waiting".to_string()), ),
+    ).await.unwrap_err();
+
+    // unsupported: enum param as String
+    client.query_required_single::<bool, (String, )>(
+        enum_query, &("waiting".to_string(), ),
+    ).await.unwrap_err();
+
+    // enum param as &str
+    assert_eq!(
+        client.query_required_single::<bool, (&'_ str, )>(
+            enum_query, &("waiting", ),
+        ).await.unwrap(),
+        true
+    );
+
     Ok(())
 }
 

diff --git a/edgedb-tokio/tests/func/dbschema/test.esdl b/edgedb-tokio/tests/func/dbschema/test.esdl
@@ -1,4 +1,6 @@
 module test {
+    scalar type State extending enum<'done', 'waiting', 'blocked'>;
+
     type Counter {
         required property name -> str {
             constraint std::exclusive;