Panache: support generating count queries with more than one column

Using a subselect. Only useful for HR until they support count queries,
but it was easy enough to fix

extensions/panache/panache-hibernate-common/runtime/src/main/java/io/quarkus/panache/hibernate/common/runtime/CountParserVisitor.java

@@ -43,8 +43,7 @@ public String visitSelectClause(SelectClauseContext ctx) {
                 sb.append(" count(");
                 ctx.DISTINCT().accept(this);
                 if (ctx.selectionList().children.size() != 1) {
-                    // FIXME: error message should include query
-                    throw new RuntimeException("Cannot count on more than one column");
+                    throw new RequiresSubqueryException();
                 }
                 ctx.selectionList().children.get(0).accept(this);
                 sb.append(" )");

extensions/panache/panache-hibernate-common/runtime/src/main/java/io/quarkus/panache/hibernate/common/runtime/PanacheJpaUtil.java

@@ -97,9 +97,17 @@ public static String getCountQueryUsingParser(String query) {
         CommonTokenStream tokens = new CommonTokenStream(lexer);
         HqlParser parser = new HqlParser(tokens);
         SelectStatementContext statement = parser.selectStatement();
-        CountParserVisitor visitor = new CountParserVisitor();
-        statement.accept(visitor);
-        return visitor.result();
+        try {
+            CountParserVisitor visitor = new CountParserVisitor();
+            statement.accept(visitor);
+            return visitor.result();
+        } catch (RequiresSubqueryException x) {
+            // no luck
+            SubQueryAliasParserVisitor visitor = new SubQueryAliasParserVisitor();
+            statement.accept(visitor);
+            String ret = visitor.result();
+            return "select count( * ) from ( " + ret + " )";
+        }
     }
 
     public static String getEntityName(Class<?> entityClass) {

extensions/panache/panache-hibernate-common/runtime/src/main/java/io/quarkus/panache/hibernate/common/runtime/RequiresSubqueryException.java

@@ -0,0 +1,5 @@
+package io.quarkus.panache.hibernate.common.runtime;
+
+public class RequiresSubqueryException extends RuntimeException {
+
+}

extensions/panache/panache-hibernate-common/runtime/src/main/java/io/quarkus/panache/hibernate/common/runtime/SubQueryAliasParserVisitor.java

@@ -0,0 +1,85 @@
+package io.quarkus.panache.hibernate.common.runtime;
+
+import org.antlr.v4.runtime.tree.TerminalNode;
+import org.hibernate.grammars.hql.HqlParser.JoinContext;
+import org.hibernate.grammars.hql.HqlParser.QueryOrderContext;
+import org.hibernate.grammars.hql.HqlParser.SelectionContext;
+import org.hibernate.grammars.hql.HqlParser.SimpleQueryGroupContext;
+import org.hibernate.grammars.hql.HqlParserBaseVisitor;
+
+public class SubQueryAliasParserVisitor extends HqlParserBaseVisitor<String> {
+
+    private int inSimpleQueryGroup;
+    private StringBuilder sb = new StringBuilder();
+    private int counter;
+
+    @Override
+    public String visitSimpleQueryGroup(SimpleQueryGroupContext ctx) {
+        inSimpleQueryGroup++;
+        try {
+            return super.visitSimpleQueryGroup(ctx);
+        } finally {
+            inSimpleQueryGroup--;
+        }
+    }
+
+    @Override
+    public String visitSelection(SelectionContext ctx) {
+        super.visitSelection(ctx);
+        if (inSimpleQueryGroup == 1) {
+            if (ctx.variable() == null) {
+                sb.append(" as __v" + counter++);
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public String visitJoin(JoinContext ctx) {
+        if (inSimpleQueryGroup == 1 && ctx.FETCH() != null) {
+            // ignore fetch joins for main query
+            return null;
+        }
+        return super.visitJoin(ctx);
+    }
+
+    @Override
+    public String visitQueryOrder(QueryOrderContext ctx) {
+        if (inSimpleQueryGroup == 1) {
+            // ignore order/limit/offset for main query
+            return null;
+        }
+        return super.visitQueryOrder(ctx);
+    }
+
+    @Override
+    public String visitTerminal(TerminalNode node) {
+        append(node.getText());
+        return null;
+    }
+
+    @Override
+    protected String defaultResult() {
+        return null;
+    }
+
+    @Override
+    protected String aggregateResult(String aggregate, String nextResult) {
+        if (nextResult != null) {
+            append(nextResult);
+        }
+        return null;
+    }
+
+    private void append(String nextResult) {
+        // don't add space at start, or around dots, commas
+        if (!sb.isEmpty() && sb.charAt(sb.length() - 1) != '.' && !nextResult.equals(".") && !nextResult.equals(",")) {
+            sb.append(" ");
+        }
+        sb.append(nextResult);
+    }
+
+    public String result() {
+        return sb.toString();
+    }
+}

extensions/panache/panache-hibernate-common/runtime/src/test/java/io/quarkus/panache/hibernate/common/runtime/CountTest.java

@@ -13,8 +13,10 @@ public void testParser() {
         // one column distinct
         assertCountQueryUsingParser("select count( distinct foo ) from bar", "select distinct foo from bar");
         // two columns distinct
-        Assertions.assertThrows(RuntimeException.class,
-                () -> assertCountQueryUsingParser("XX", "select distinct foo,gee from bar"));
+        assertCountQueryUsingParser("select count( * ) from ( select distinct foo as __v0, gee as __v1 from bar )",
+                "select distinct foo,gee from bar");
+        assertCountQueryUsingParser("select count( * ) from ( select distinct foo as __v0, gee as g from bar )",
+                "select distinct foo,gee as g from bar");
         // nested order by not touched
         assertCountQueryUsingParser("select count( * ) from ( from entity order by id )",
                 "select foo from (from entity order by id) order by foo, bar ASC");