compiler: Handle subqueries in SELECT statements (#1113)

Author: kyleconroy

internal/compiler/output_columns.go

@@ -3,6 +3,7 @@ package compiler
 import (
 	"errors"
 	"fmt"
+
 	"github.com/kyleconroy/sqlc/internal/sql/catalog"
 
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
@@ -202,6 +203,16 @@ func outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, error) {
 			switch n.SubLinkType {
 			case ast.EXISTS_SUBLINK:
 				cols = append(cols, &Column{Name: name, DataType: "bool", NotNull: true})
+			case ast.EXPR_SUBLINK:
+				subcols, err := outputColumns(qc, n.Subselect)
+				if err != nil {
+					return nil, err
+				}
+				first := subcols[0]
+				if res.Name != nil {
+					first.Name = *res.Name
+				}
+				cols = append(cols, first)
 			default:
 				cols = append(cols, &Column{Name: name, DataType: "any", NotNull: false})
 			}

internal/endtoend/testdata/select_nested_count/postgresql/go/db.go

@@ -0,0 +1,19 @@
+CREATE TABLE authors (
+  id   BIGSERIAL PRIMARY KEY,
+  name text      NOT NULL,
+  bio  text
+);
+
+CREATE TABLE books (
+  id BIGSERIAL PRIMARY KEY,
+  author_id BIGSERIAL NOT NULL
+  	REFERENCES authors(id),
+  title text NOT NULL
+);
+
+-- name: GetAuthorsWithBooksCount :many
+SELECT *, (
+  SELECT COUNT(id) FROM books
+  WHERE books.author_id = id
+) AS books_count
+FROM authors;

internal/endtoend/testdata/select_nested_count/postgresql/go/models.go

@@ -0,0 +1,11 @@
+{
+  "version": "1",
+  "packages": [
+    {
+      "path": "go",
+      "name": "querytest",
+      "schema": "query.sql",
+      "queries": "query.sql"
+    }
+  ]
+}

internal/endtoend/testdata/select_nested_count/postgresql/go/query.sql.go

@@ -38,7 +38,7 @@ func convertSubLinkType(t pg.SubLinkType) (ast.SubLinkType, error) {
 	case pg.SubLinkType_ROWCOMPARE_SUBLINK:
 		return ast.ROWCOMPARE_SUBLINK, nil
 	case pg.SubLinkType_EXPR_SUBLINK:
-		return ast.EXISTS_SUBLINK, nil
+		return ast.EXPR_SUBLINK, nil
 	case pg.SubLinkType_MULTIEXPR_SUBLINK:
 		return ast.MULTIEXPR_SUBLINK, nil
 	case pg.SubLinkType_ARRAY_SUBLINK:
@@ -2775,10 +2775,10 @@ func convertVacuumStmt(n *pg.VacuumStmt) *ast.VacuumStmt {
 		return nil
 	}
 	return &ast.VacuumStmt{
-		// FIXME: The VacuumStmt node has changed quite a bit
-		// Options:  n.Options
-		// Relation: convertRangeVar(n.Relation),
-		// VaCols:   convertSlice(n.VaCols),
+	// FIXME: The VacuumStmt node has changed quite a bit
+	// Options:  n.Options
+	// Relation: convertRangeVar(n.Relation),
+	// VaCols:   convertSlice(n.VaCols),
 	}
 }