Fix retrieving multiple columns from the same response (#4678)

sdk/src/api/method/query.rs

@@ -916,6 +916,31 @@ mod tests {
 		assert_eq!(value.into_inner(), CoreValue::from(article.title));
 	}
 
+	#[test]
+	fn take_key_multi() {
+		let article = Article {
+			title: "Lorem Ipsum".to_owned(),
+			body: "Lorem Ipsum Lorem Ipsum".to_owned(),
+		};
+		let value = to_value(article.clone()).unwrap();
+
+		let mut response = Response {
+			results: to_map(vec![Ok(value.clone().into_inner())]),
+			..Response::new()
+		};
+		let title: Vec<String> = response.take("title").unwrap();
+		assert_eq!(title, vec![article.title.clone()]);
+		let body: Vec<String> = response.take("body").unwrap();
+		assert_eq!(body, vec![article.body]);
+
+		let mut response = Response {
+			results: to_map(vec![Ok(value.clone().into_inner())]),
+			..Response::new()
+		};
+		let vec: Vec<String> = response.take("title").unwrap();
+		assert_eq!(vec, vec![article.title]);
+	}
+
 	#[test]
 	fn take_partial_records() {
 		let mut response = Response {

sdk/src/api/opt/query.rs

@@ -377,34 +377,37 @@ where
 {
 	fn query_result(self, response: &mut QueryResponse) -> Result<Vec<T>> {
 		let (index, key) = self;
-		let mut response = match response.results.get_mut(&index) {
+		match response.results.get_mut(&index) {
 			Some((_, result)) => match result {
 				Ok(val) => match val {
-					CoreValue::Array(vec) => mem::take(&mut vec.0),
+					CoreValue::Array(vec) => {
+						let mut responses = Vec::with_capacity(vec.len());
+						for value in vec.iter_mut() {
+							if let CoreValue::Object(object) = value {
+								if let Some(value) = object.remove(key) {
+									responses.push(value);
+								}
+							}
+						}
+						from_core_value(responses.into()).map_err(Into::into)
+					}
 					val => {
-						let val = mem::take(val);
-						vec![val]
+						if let CoreValue::Object(object) = val {
+							if let Some(value) = object.remove(key) {
+								return from_core_value(vec![value].into()).map_err(Into::into);
+							}
+						}
+						Ok(vec![])
 					}
 				},
 				Err(error) => {
 					let error = mem::replace(error, Error::ConnectionUninitialised.into());
 					response.results.swap_remove(&index);
-					return Err(error);
+					Err(error)
 				}
 			},
-			None => {
-				return Ok(vec![]);
-			}
-		};
-		let mut vec = Vec::with_capacity(response.len());
-		for value in response.iter_mut() {
-			if let CoreValue::Object(object) = value {
-				if let Some(value) = object.remove(key) {
-					vec.push(value);
-				}
-			}
+			None => Ok(vec![]),
 		}
-		from_core_value(vec.into()).map_err(Into::into)
 	}
 
 	fn stats(&self, response: &QueryResponse) -> Option<Stats> {

sdk/tests/api/mod.rs

@@ -580,7 +580,8 @@ async fn insert_relation_table() {
 	let (permit, db) = new_db().await;
 	db.use_ns(NS).use_db(Ulid::new().to_string()).await.unwrap();
 	drop(permit);
-	let tmp: Result<Vec<ApiRecordId>, _> = db.insert("likes").relation("{}".parse::<Value>().unwrap()).await;
+	let tmp: Result<Vec<ApiRecordId>, _> =
+		db.insert("likes").relation("{}".parse::<Value>().unwrap()).await;
 	tmp.unwrap_err();
 	let val = "{in: person:a, out: thing:a}".parse::<Value>().unwrap();
 	let _: Vec<ApiRecordId> = db.insert("likes").relation(val).await.unwrap();
@@ -1453,3 +1454,23 @@ async fn run() {
 	let tmp: i32 = db.run("fn::baz").await.unwrap();
 	assert_eq!(tmp, 7);
 }
+
+#[test_log::test(tokio::test)]
+async fn multi_take() {
+	let (permit, db) = new_db().await;
+	db.use_ns(NS).use_db(Ulid::new().to_string()).await.unwrap();
+	drop(permit);
+
+	db.query("INSERT INTO user {name: 'John', address: 'USA'};").await.unwrap();
+	db.query("INSERT INTO user {name: 'Adam', address: 'UK'};").await.unwrap();
+
+	let mut response = db.query("SELECT * FROM user").await.unwrap();
+
+	let mut names: Vec<String> = response.take("name").unwrap();
+	names.sort();
+	assert_eq!(names, vec!["Adam".to_owned(), "John".to_owned()]);
+
+	let mut addresses: Vec<String> = response.take("address").unwrap();
+	addresses.sort();
+	assert_eq!(addresses, vec!["UK".to_owned(), "USA".to_owned()]);
+}