Repository: tomarrell/lbadd (go)

Issue:
engine: projection should support alias
`evaluateProject` does not handle `alias` currently. There is a TODO in the code.

Pull request:
Fix PRs #203 and #205
Implement changes from #203 and #205, which were added on the wrong repo by mistake.

Closes #193

## Definition of done
- [ ] Code correctness
- [ ] Documentation
- [ ] Test cases


Patch:
diff --git a/internal/engine/error.go b/internal/engine/error.go
index ab4e515..0e88099 100644
--- a/internal/engine/error.go
+++ b/internal/engine/error.go
@@ -41,3 +41,9 @@ func ErrUncomparable(t types.Type) Error {
 func ErrUnimplemented(what interface{}) Error {
 	return Error(fmt.Sprintf("'%v' is not implemented", what))
 }
+
+// ErrNoSuchColumn returns an error indicating that a requested column is not
+// contained in the current result table.
+func ErrNoSuchColumn(name string) Error {
+	return Error(fmt.Sprintf("no column with name or alias '%s'", name))
+}
diff --git a/internal/engine/evaluate.go b/internal/engine/evaluate.go
index 9378b47..35dcd7e 100644
--- a/internal/engine/evaluate.go
+++ b/internal/engine/evaluate.go
@@ -50,6 +50,7 @@ func (e Engine) evaluateProjection(ctx ExecutionContext, proj command.Project) (
 	}
 
 	var expectedColumnNames []string
+	aliases := make(map[string]string)
 	for _, col := range proj.Cols {
 		// evaluate the column name
 		colNameExpr, err := e.evaluateExpression(ctx, col.Column)
@@ -59,16 +60,37 @@ func (e Engine) evaluateProjection(ctx ExecutionContext, proj command.Project) (
 		var colName string
 		if colNameExpr.Is(types.String) {
 			colName = colNameExpr.(types.StringValue).Value
+		} else {
+			casted, err := types.String.Cast(colNameExpr)
+			if err != nil {
+				return Table{}, fmt.Errorf("cannot cast %v to %v: %w", colNameExpr.Type(), types.String, err)
+			}
+			colName = casted.(types.StringValue).Value
 		}
 		if col.Table != "" {
 			colName = col.Table + "." + colName
 		}
-
-		// #193: support alias
+		if col.Alias != "" {
+			aliases[colName] = col.Alias
+		}
 
 		expectedColumnNames = append(expectedColumnNames, colName)
 	}
 
+	// check if the table actually has all expected columns
+	for _, expectedCol := range expectedColumnNames {
+		if !origin.HasColumn(expectedCol) {
+			return Table{}, ErrNoSuchColumn(expectedCol)
+		}
+	}
+
+	// apply aliases
+	for i, col := range origin.Cols {
+		if alias, ok := aliases[col.QualifiedName]; ok {
+			origin.Cols[i].Alias = alias
+		}
+	}
+
 	sort.Strings(expectedColumnNames)
 
 	var toRemove []string
diff --git a/internal/engine/table.go b/internal/engine/table.go
index 7d88bed..de8afb1 100644
--- a/internal/engine/table.go
+++ b/internal/engine/table.go
@@ -54,6 +54,17 @@ func (t Table) RemoveColumnByQualifiedName(qualifiedName string) Table {
 	return t
 }
 
+// HasColumn inspects the table's columns and determines whether the table has
+// any column, that has the given name as qualified name OR as alias.
+func (t Table) HasColumn(qualifiedNameOrAlias string) bool {
+	for _, col := range t.Cols {
+		if col.QualifiedName == qualifiedNameOrAlias || col.Alias == qualifiedNameOrAlias {
+			return true
+		}
+	}
+	return false
+}
+
 // RemoveColumn works on a copy of the table, and removes the column with the
 // given index from the copy. After removal, the copy is returned.
 func (t Table) RemoveColumn(index int) Table {


Test patch:
diff --git a/internal/engine/evaluate_test.go b/internal/engine/evaluate_test.go
index 7d4ab7b..00bc06f 100644
--- a/internal/engine/evaluate_test.go
+++ b/internal/engine/evaluate_test.go
@@ -23,38 +23,134 @@ func TestFullTableScan(t *testing.T) {
 	assert.Equal(Table{}, result)
 }
 
-func TestProjection(t *testing.T) {
-	assert := assert.New(t)
-
-	e := createEngineOnEmptyDatabase(t)
-	result, err := e.Evaluate(command.Project{
-		Cols: []command.Column{
-			{
-				Column: command.LiteralExpr{Value: "column2"},
+func TestEngine_evaluateProjection(t *testing.T) {
+	tests := []struct {
+		name    string
+		ctx     ExecutionContext
+		proj    command.Project
+		want    Table
+		wantErr string
+	}{
+		{
+			"empty",
+			newEmptyExecutionContext(),
+			command.Project{
+				Cols: []command.Column{},
+				Input: command.Values{
+					Values: [][]command.Expr{
+						{command.LiteralExpr{Value: "hello"}, command.LiteralExpr{Value: "world"}, command.ConstantBooleanExpr{Value: true}},
+						{command.LiteralExpr{Value: "foo"}, command.LiteralExpr{Value: "bar"}, command.ConstantBooleanExpr{Value: false}},
+					},
+				},
 			},
-		},
-		Input: command.Values{
-			Values: [][]command.Expr{
-				{command.LiteralExpr{Value: "hello"}, command.LiteralExpr{Value: "world"}, command.ConstantBooleanExpr{Value: true}},
-				{command.LiteralExpr{Value: "foo"}, command.LiteralExpr{Value: "bar"}, command.ConstantBooleanExpr{Value: false}},
+			Table{
+				Cols: []Col{},
+				Rows: []Row{},
 			},
+			"",
 		},
-	})
-	assert.NoError(err)
-	assert.Equal(Table{
-		Cols: []Col{
-			{
-				QualifiedName: "column2",
-				Type:          types.String,
+		{
+			"simple",
+			newEmptyExecutionContext(),
+			command.Project{
+				Cols: []command.Column{
+					{
+						Column: command.LiteralExpr{Value: "column2"},
+					},
+				},
+				Input: command.Values{
+					Values: [][]command.Expr{
+						{command.LiteralExpr{Value: "hello"}, command.LiteralExpr{Value: "world"}, command.ConstantBooleanExpr{Value: true}},
+						{command.LiteralExpr{Value: "foo"}, command.LiteralExpr{Value: "bar"}, command.ConstantBooleanExpr{Value: false}},
+					},
+				},
+			},
+			Table{
+				Cols: []Col{
+					{
+						QualifiedName: "column2",
+						Type:          types.String,
+					},
+				},
+				Rows: []Row{
+					{
+						Values: []types.Value{types.NewString("world")},
+					},
+					{
+						Values: []types.Value{types.NewString("bar")},
+					},
+				},
 			},
+			"",
 		},
-		Rows: []Row{
-			{
-				Values: []types.Value{types.NewString("world")},
+		{
+			"simple with alias",
+			newEmptyExecutionContext(),
+			command.Project{
+				Cols: []command.Column{
+					{
+						Column: command.LiteralExpr{Value: "column2"},
+						Alias:  "foo",
+					},
+				},
+				Input: command.Values{
+					Values: [][]command.Expr{
+						{command.LiteralExpr{Value: "hello"}, command.LiteralExpr{Value: "world"}, command.ConstantBooleanExpr{Value: true}},
+						{command.LiteralExpr{Value: "foo"}, command.LiteralExpr{Value: "bar"}, command.ConstantBooleanExpr{Value: false}},
+					},
+				},
 			},
-			{
-				Values: []types.Value{types.NewString("bar")},
+			Table{
+				Cols: []Col{
+					{
+						QualifiedName: "column2",
+						Alias:         "foo",
+						Type:          types.String,
+					},
+				},
+				Rows: []Row{
+					{
+						Values: []types.Value{types.NewString("world")},
+					},
+					{
+						Values: []types.Value{types.NewString("bar")},
+					},
+				},
 			},
+			"",
 		},
-	}, result)
+		{
+			"missing column",
+			newEmptyExecutionContext(),
+			command.Project{
+				Cols: []command.Column{
+					{
+						Column: command.LiteralExpr{Value: "foo"},
+					},
+				},
+				Input: command.Values{
+					Values: [][]command.Expr{
+						{command.LiteralExpr{Value: "hello"}},
+						{command.LiteralExpr{Value: "foo"}},
+					},
+				},
+			},
+			Table{},
+			"no column with name or alias 'foo'",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			e := createEngineOnEmptyDatabase(t)
+			got, err := e.evaluateProjection(tt.ctx, tt.proj)
+			if tt.wantErr != "" {
+				assert.EqualError(err, tt.wantErr)
+			} else {
+				assert.NoError(err)
+			}
+			assert.Equal(tt.want, got)
+		})
+	}
 }