From 0e9903349c0468a84db63761d00f4d5d08296679 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyle@conroy.org>
Date: Mon, 18 Nov 2019 12:46:56 -0800
Subject: [PATCH] parser: Add support for mathmatical operators

Also, don't leave out expressions we haven't added. Instead, generate a
column with an "any" type.
---
 internal/dinosql/parser.go     |  9 ++++++++-
 internal/dinosql/query_test.go | 13 +++++++++++++
 internal/postgres/types.go     | 25 +++++++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/internal/dinosql/parser.go b/internal/dinosql/parser.go
index ae157f9519..a5a12a7c8d 100644
--- a/internal/dinosql/parser.go
+++ b/internal/dinosql/parser.go
@@ -379,6 +379,7 @@ func parseQuery(c core.Catalog, stmt nodes.Node, source string) (*Query, error)
 	default:
 		return nil, nil
 	}
+
 	rawSQL, err := pluckQuery(source, raw)
 	if err != nil {
 		return nil, err
@@ -562,9 +563,15 @@ func outputColumns(c core.Catalog, node nodes.Node) ([]core.Column, error) {
 			if res.Name != nil {
 				name = *res.Name
 			}
-			if postgres.IsComparisonOperator(join(n.Name, "")) {
+			switch {
+			case postgres.IsComparisonOperator(join(n.Name, "")):
 				// TODO: Generate a name for these operations
 				cols = append(cols, core.Column{Name: name, DataType: "bool", NotNull: true})
+			case postgres.IsMathematicalOperator(join(n.Name, "")):
+				// TODO: Generate correct numeric type
+				cols = append(cols, core.Column{Name: name, DataType: "pg_catalog.int4", NotNull: true})
+			default:
+				cols = append(cols, core.Column{Name: name, DataType: "any", NotNull: false})
 			}
 
 		case nodes.CoalesceExpr:
diff --git a/internal/dinosql/query_test.go b/internal/dinosql/query_test.go
index 59d865868c..06bdd63534 100644
--- a/internal/dinosql/query_test.go
+++ b/internal/dinosql/query_test.go
@@ -690,6 +690,19 @@ func TestQueries(t *testing.T) {
 				},
 			},
 		},
+		{
+			"mathmatical-operator",
+			`
+			CREATE TABLE foo (num integer not null);
+			SELECT *, num / 1024 as division FROM foo;
+			`,
+			Query{
+				Columns: []core.Column{
+					{Name: "num", DataType: "pg_catalog.int4", NotNull: true, Table: public("foo")},
+					{Name: "division", DataType: "pg_catalog.int4", NotNull: true},
+				},
+			},
+		},
 	} {
 		test := tc
 		t.Run(test.name, func(t *testing.T) {
diff --git a/internal/postgres/types.go b/internal/postgres/types.go
index df6793dea0..e1bb317d65 100644
--- a/internal/postgres/types.go
+++ b/internal/postgres/types.go
@@ -14,3 +14,28 @@ func IsComparisonOperator(s string) bool {
 	}
 	return true
 }
+
+func IsMathematicalOperator(s string) bool {
+	switch s {
+	case "+":
+	case "-":
+	case "*":
+	case "/":
+	case "%":
+	case "^":
+	case "|/":
+	case "||/":
+	case "!":
+	case "!!":
+	case "@":
+	case "&":
+	case "|":
+	case "#":
+	case "~":
+	case "<<":
+	case ">>":
+	default:
+		return false
+	}
+	return true
+}