From f7bc7c21316c2cfdc55880276d0d617afd8e382c Mon Sep 17 00:00:00 2001
From: George Georgiev <george@harness.io>
Date: Tue, 8 Jun 2021 13:23:37 -0700
Subject: [PATCH] Add projection generation

---
 src/compiler/compiler.go      | 9 +++++++--
 src/compiler/compiler_test.go | 9 ++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/src/compiler/compiler.go b/src/compiler/compiler.go
index 962272a..2a4f5ec 100644
--- a/src/compiler/compiler.go
+++ b/src/compiler/compiler.go
@@ -11,7 +11,7 @@ type Compiler struct {
 
 const createTemplate = `
   public static %sQuery%s create(HPersistence persistence) {
-    return new QueryImpl(persistence.createQuery(%s.class));
+    return new QueryImpl(persistence.createQuery(%s.class)%s);
   }`
 
 const interfaceTemplate = `
@@ -87,7 +87,12 @@ func (compiler *Compiler) Generate(query *dom.Query) string {
 	var name = query.Name
 	var collectionName = compiler.collectionName(query.Collection)
 
-	createMethod := fmt.Sprintf(createTemplate, name, strings.Title(query.Filters[0].FieldName), collectionName)
+	var projections strings.Builder
+	for _, field := range query.ProjectFields {
+		projections.WriteString(fmt.Sprintf("\n                                    .project(%sKeys.%s, true)", collectionName, field))
+	}
+
+	createMethod := fmt.Sprintf(createTemplate, name, strings.Title(query.Filters[0].FieldName), collectionName, projections.String())
 
 	var interfaces strings.Builder
 	var interfaceNames strings.Builder
diff --git a/src/compiler/compiler_test.go b/src/compiler/compiler_test.go
index fa24054..ce4e4a0 100644
--- a/src/compiler/compiler_test.go
+++ b/src/compiler/compiler_test.go
@@ -80,6 +80,7 @@ public class DelegateTaskSelectQuery implements PersistentQuery {
 			{FieldType: "String", FieldName: "accountId"},
 			{FieldType: "String", FieldName: "uuid", Operation: dom.In},
 		},
+		ProjectFields: []string{"foo", "bar"},
 	}
 	expected1 := `package io.harness.beans;
 
@@ -91,7 +92,9 @@ import org.mongodb.morphia.query.Query;
 
 public class DelegateTaskSelectQuery implements PersistentQuery {
   public static SelectQueryAccountId create(HPersistence persistence) {
-    return new QueryImpl(persistence.createQuery(DelegateTask.class));
+    return new QueryImpl(persistence.createQuery(DelegateTask.class)
+                                    .project(DelegateTaskKeys.foo, true)
+                                    .project(DelegateTaskKeys.bar, true));
   }
 
   public interface SelectQueryAccountId {
@@ -125,6 +128,10 @@ public class DelegateTaskSelectQuery implements PersistentQuery {
       return query;
     }
   }
+
+  List<String> queryCanonicalForms() {
+    return ImmutableList.<String>builder().build();
+  }
 }
 `
 	result = compiler.Generate(&query1)