From 5d2193bd6398a104631e13be8a835ba3c7d122ec Mon Sep 17 00:00:00 2001
From: Richard Newman <rnewman@twinql.com>
Date: Thu, 13 Apr 2017 11:45:38 -0700
Subject: [PATCH] Part 2: include necessary 'with' variables in SQL projection
 list.

The test produces projection elements for `:with`, even though there are
no aggregates in the query. This test will need to be adjusted when we
optimize this away!
---
 query-projector/src/lib.rs          | 15 +++++++++++++++
 query-translator/tests/translate.rs | 17 +++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/query-projector/src/lib.rs b/query-projector/src/lib.rs
index bc5d0d2a..0e9e19ed 100644
--- a/query-projector/src/lib.rs
+++ b/query-projector/src/lib.rs
@@ -199,6 +199,7 @@ fn project_elements<'a, I: IntoIterator<Item = &'a Element>>(
     let mut cols = Vec::with_capacity(count);
     let mut i: i32 = 0;
     let mut templates = vec![];
+    let mut with = query.with.clone();
 
     for e in elements {
         match e {
@@ -206,6 +207,8 @@ fn project_elements<'a, I: IntoIterator<Item = &'a Element>>(
             // into the SQL projection, aliased to the name of the variable,
             // and we push an annotated index into the projector.
             &Element::Variable(ref var) => {
+                // If we're projecting this, we don't need it in :with.
+                with.remove(var);
 
                 let (column, name) = candidate_column(query, var);
                 cols.push(ProjectedColumn(column, name));
@@ -225,6 +228,18 @@ fn project_elements<'a, I: IntoIterator<Item = &'a Element>>(
         }
     }
 
+    for var in with {
+        // We need to collect these into the SQL column list, but they don't affect projection.
+        // If a variable is of a non-fixed type, also project the type tag column, so we don't
+        // accidentally unify across types when considering uniqueness!
+        let (column, name) = candidate_column(query, &var);
+        cols.push(ProjectedColumn(column, name));
+        if query.cc.known_type(&var).is_none() {
+            let (type_column, type_name) = candidate_type_column(query, &var);
+            cols.push(ProjectedColumn(type_column, type_name));
+        }
+    }
+
     (Projection::Columns(cols), templates)
 }
 
diff --git a/query-translator/tests/translate.rs b/query-translator/tests/translate.rs
index 02445727..e0af84e5 100644
--- a/query-translator/tests/translate.rs
+++ b/query-translator/tests/translate.rs
@@ -353,3 +353,20 @@ fn test_complex_or_join_type_projection() {
                     LIMIT 1");
     assert_eq!(args, vec![]);
 }
+
+#[test]
+fn test_with_without_aggregate() {
+    let schema = prepopulated_schema();
+
+    // Known type.
+    let input = r#"[:find ?x :with ?y :where [?x :foo/bar ?y]]"#;
+    let SQLQuery { sql, args } = translate(&schema, input, None);
+    assert_eq!(sql, "SELECT DISTINCT `datoms00`.e AS `?x`, `datoms00`.v AS `?y` FROM `datoms` AS `datoms00` WHERE `datoms00`.a = 99");
+    assert_eq!(args, vec![]);
+
+    // Unknown type.
+    let input = r#"[:find ?x :with ?y :where [?x _ ?y]]"#;
+    let SQLQuery { sql, args } = translate(&schema, input, None);
+    assert_eq!(sql, "SELECT DISTINCT `all_datoms00`.e AS `?x`, `all_datoms00`.v AS `?y`, `all_datoms00`.value_type_tag AS `?y_value_type_tag` FROM `all_datoms` AS `all_datoms00`");
+    assert_eq!(args, vec![]);
+}
\ No newline at end of file