Parse and handle aliased pull attributes.

query-parser/src/parse.rs

@@ -67,6 +67,7 @@ use self::mentat_query::{
     Order,
     OrJoin,
     OrWhereClause,
+    NamedPullAttribute,
     NotJoin,
     Pattern,
     PatternNonValuePlace,
@@ -191,6 +192,7 @@ def_parser!(Query, order, Order, {
 def_matches_plain_symbol!(Query, the, "the");
 def_matches_plain_symbol!(Query, pull, "pull");
 def_matches_plain_symbol!(Query, wildcard, "*");
+def_matches_keyword!(Query, alias_as, "as");
 
 pub struct Where<'a>(std::marker::PhantomData<&'a ()>);
 
@@ -303,11 +305,28 @@ def_parser!(Query, aggregate, Aggregate, {
          })
 });
 
+def_parser!(Query, pull_concrete_attribute_ident, PullConcreteAttribute, {
+    forward_keyword().map(|k| PullConcreteAttribute::Ident(::std::rc::Rc::new(k.clone())))
+});
+
+def_parser!(Query, pull_aliased_attribute, PullAttributeSpec, {
+    vector().of_exactly(
+        (Query::pull_concrete_attribute_ident()
+            .skip(Query::alias_as()),
+        forward_keyword().map(|alias| Some(::std::rc::Rc::new(alias.clone()))))
+            .map(|(attribute, alias)|
+                PullAttributeSpec::Attribute(
+                    NamedPullAttribute { attribute, alias })))
+});
+
 def_parser!(Query, pull_concrete_attribute, PullAttributeSpec, {
-    forward_keyword().map(|k|
+    Query::pull_concrete_attribute_ident()
-        PullAttributeSpec::Attribute(
+        .map(|attribute|
-            PullConcreteAttribute::Ident(
+            PullAttributeSpec::Attribute(
-                ::std::rc::Rc::new(k.clone()))))
+                NamedPullAttribute {
+                    attribute,
+                    alias: None,
+                }))
 });
 
 def_parser!(Query, pull_wildcard_attribute, PullAttributeSpec, {
@@ -316,6 +335,7 @@ def_parser!(Query, pull_wildcard_attribute, PullAttributeSpec, {
 
 def_parser!(Query, pull_attribute, PullAttributeSpec, {
     choice([
+        try(Query::pull_aliased_attribute()),
         try(Query::pull_concrete_attribute()),
         try(Query::pull_wildcard_attribute()),
         // TODO: reversed keywords, entids (with aliases, presumably…).
@@ -1205,23 +1225,27 @@ mod test {
 
         let foo_bar = ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "bar"));
         let foo_baz = ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "baz"));
+        let foo_horse = ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "horse"));
         assert_edn_parses_to!(Query::pull_concrete_attribute,
                               ":foo/bar",
                               PullAttributeSpec::Attribute(
-                                  PullConcreteAttribute::Ident(foo_bar.clone())));
+                                  PullConcreteAttribute::Ident(foo_bar.clone()).into()));
         assert_edn_parses_to!(Query::pull_attribute,
                               ":foo/bar",
                               PullAttributeSpec::Attribute(
-                                  PullConcreteAttribute::Ident(foo_bar.clone())));
+                                  PullConcreteAttribute::Ident(foo_bar.clone()).into()));
         assert_edn_parses_to!(Find::elem,
-                              "(pull ?v [:foo/bar :foo/baz])",
+                              "(pull ?v [[:foo/bar :as :foo/horse] :foo/baz])",
                               Element::Pull(Pull {
                                   var: Variable::from_valid_name("?v"),
                                   patterns: vec![
                                       PullAttributeSpec::Attribute(
-                                          PullConcreteAttribute::Ident(foo_bar.clone())),
+                                          NamedPullAttribute {
+                                              attribute: PullConcreteAttribute::Ident(foo_bar.clone()),
+                                              alias: Some(foo_horse),
+                                          }),
                                       PullAttributeSpec::Attribute(
-                                          PullConcreteAttribute::Ident(foo_baz.clone())),
+                                          PullConcreteAttribute::Ident(foo_baz.clone()).into()),
                                   ],
                               }));
         assert_parse_failure_contains!(Find::elem,
@@ -1242,7 +1266,7 @@ mod test {
                             PullAttributeSpec::Attribute(
                                 PullConcreteAttribute::Ident(
                                     ::std::rc::Rc::new(edn::NamespacedKeyword::new("foo", "bar"))
-                                )
+                                ).into()
                             ),
                         ] })]),
                 where_clauses: vec![

query-pull/src/lib.rs

@@ -91,6 +91,7 @@ use mentat_core::{
 use mentat_db::cache;
 
 use mentat_query::{
+    NamedPullAttribute,
     PullAttributeSpec,
     PullConcreteAttribute,
 };
@@ -110,7 +111,7 @@ pub fn pull_attributes_for_entity<A>(schema: &Schema,
                                      attributes: A) -> Result<StructuredMap>
     where A: IntoIterator<Item=Entid> {
     let attrs = attributes.into_iter()
-                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e)))
+                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e).into()))
                           .collect();
     Puller::prepare(schema, attrs)?
         .pull(schema, db, once(entity))
@@ -130,7 +131,7 @@ pub fn pull_attributes_for_entities<E, A>(schema: &Schema,
     where E: IntoIterator<Item=Entid>,
           A: IntoIterator<Item=Entid> {
     let attrs = attributes.into_iter()
-                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e)))
+                          .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e).into()))
                           .collect();
     Puller::prepare(schema, attrs)?
         .pull(schema, db, entities)
@@ -148,7 +149,7 @@ impl Puller {
     pub fn prepare_simple_attributes(schema: &Schema, attributes: Vec<Entid>) -> Result<Puller> {
         Puller::prepare(schema,
                         attributes.into_iter()
-                                  .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e)))
+                                  .map(|e| PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(e).into()))
                                   .collect())
     }
 
@@ -175,16 +176,27 @@ impl Puller {
                     }
                     break;
                 },
-                &PullAttributeSpec::Attribute(PullConcreteAttribute::Ident(ref i)) => {
+                &PullAttributeSpec::Attribute(NamedPullAttribute {
-                    if let Some(entid) = schema.get_entid(i) {
+                    ref attribute,
-                        names.insert(entid.into(), i.to_value_rc());
+                    ref alias,
-                        attrs.insert(entid.into());
+                }) => {
+                    let alias = alias.as_ref()
+                                     .map(|ref r| r.to_value_rc());
+                    match attribute {
+                        &PullConcreteAttribute::Ident(ref i) => {
+                            if let Some(entid) = schema.get_entid(i) {
+                                let name = alias.unwrap_or_else(|| i.to_value_rc());
+                                names.insert(entid.into(), name);
+                                attrs.insert(entid.into());
+                            }
+                        },
+                        &PullConcreteAttribute::Entid(ref entid) => {
+                            let name = alias.map(Ok).unwrap_or_else(|| lookup_name(entid))?;
+                            names.insert(*entid, name);
+                            attrs.insert(*entid);
+                        },
                     }
                 },
-                &PullAttributeSpec::Attribute(PullConcreteAttribute::Entid(ref entid)) => {
-                    names.insert(*entid, lookup_name(entid)?);
-                    attrs.insert(*entid);
-                },
             }
         }
 

query/src/lib.rs

@@ -499,12 +499,26 @@ pub enum PullConcreteAttribute {
     Entid(i64),
 }
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct NamedPullAttribute {
+    pub attribute: PullConcreteAttribute,
+    pub alias: Option<Rc<NamespacedKeyword>>,
+}
+
+impl From<PullConcreteAttribute> for NamedPullAttribute {
+    fn from(a: PullConcreteAttribute) -> Self {
+        NamedPullAttribute {
+            attribute: a,
+            alias: None,
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum PullAttributeSpec {
     Wildcard,
-    Attribute(PullConcreteAttribute),
+    Attribute(NamedPullAttribute),
     // PullMapSpec(Vec<…>),
-    // AttributeWithOpts(PullConcreteAttribute, …),
     // LimitedAttribute(PullConcreteAttribute, u64),  // Limit nil => Attribute instead.
     // DefaultedAttribute(PullConcreteAttribute, PullDefaultValue),
 }
@@ -522,14 +536,25 @@ impl std::fmt::Display for PullConcreteAttribute {
     }
 }
 
+impl std::fmt::Display for NamedPullAttribute {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        if let &Some(ref alias) = &self.alias {
+            write!(f, "{} :as {}", self.attribute, alias)
+        } else {
+            write!(f, "{}", self.attribute)
+        }
+    }
+}
+
+
 impl std::fmt::Display for PullAttributeSpec {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
             &PullAttributeSpec::Wildcard => {
                 write!(f, "*")
             },
-            &PullAttributeSpec::Attribute(ref a) => {
+            &PullAttributeSpec::Attribute(ref attr) => {
-                write!(f, "{}", a)
+                write!(f, "{}", attr)
             },
         }
     }

tests/pull.rs

@@ -116,7 +116,7 @@ fn test_simple_pull() {
     assert_eq!(pulled, expected);
 
     // Now test pull inside the query itself.
-    let query = r#"[:find ?hood (pull ?district [:district/name :district/region])
+    let query = r#"[:find ?hood (pull ?district [[:district/name :as :district/district] :district/region])
                     :where
                     (or [?hood :neighborhood/name "Beacon Hill"]
                         [?hood :neighborhood/name "Capitol Hill"])
@@ -128,12 +128,12 @@ fn test_simple_pull() {
                                            .expect("results");
 
     let beacon_district: Vec<(NamespacedKeyword, TypedValue)> = vec![
-        (kw!(:district/name), "Greater Duwamish".into()),
+        (kw!(:district/district), "Greater Duwamish".into()),
         (kw!(:district/region), schema.get_entid(&NamespacedKeyword::new("region", "se")).unwrap().into())
     ];
     let beacon_district: StructuredMap = beacon_district.into();
     let capitol_district: Vec<(NamespacedKeyword, TypedValue)> = vec![
-        (kw!(:district/name), "East".into()),
+        (kw!(:district/district), "East".into()),
         (kw!(:district/region), schema.get_entid(&NamespacedKeyword::new("region", "e")).unwrap().into())
     ];
     let capitol_district: StructuredMap = capitol_district.into();