Merge pull request #3070 from dolthub/angela/generated_columns

Don't prune VirtualColumnTable tables

Author: angelamayxie

enginetest/queries/generated_columns.go

@@ -1437,6 +1437,45 @@ var GeneratedColumnTests = []ScriptTest{
 			},
 		},
 	},
+	{
+		// https://github.com/dolthub/dolt/issues/8968
+		Name: "can select all columns from table with generated column",
+		SetUpScript: []string{
+			"create table t(pk int primary key, j1 json)",
+			`insert into t values (1, '{"name": "foo"}')`,
+			"alter table t add column g1 varchar(100) generated always as (json_unquote(json_extract(`j1`, '$.name')))",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "select * from t",
+				Expected: []sql.Row{{1, `{"name":"foo"}`, "foo"}},
+			},
+			{
+				Query:    "select pk, j1, g1 from t",
+				Expected: []sql.Row{{1, `{"name":"foo"}`, "foo"}},
+			},
+			{
+				Query:    "select pk, g1 from t",
+				Expected: []sql.Row{{1, "foo"}},
+			},
+			{
+				Query:    "select g1 from t",
+				Expected: []sql.Row{{"foo"}},
+			},
+			{
+				Query:    "select j1, g1 from t",
+				Expected: []sql.Row{{`{"name":"foo"}`, "foo"}},
+			},
+			{
+				Query:    "select j1 from t",
+				Expected: []sql.Row{{`{"name":"foo"}`}},
+			},
+			{
+				Query:    "select pk, j1 from t",
+				Expected: []sql.Row{{1, `{"name":"foo"}`}},
+			},
+		},
+	},
 }
 
 var BrokenGeneratedColumnTests = []ScriptTest{

sql/analyzer/symbol_resolution.go

@@ -202,36 +202,23 @@ func pruneTableCols(
 		return n, transform.SameTree, nil
 	}
 
+	// columns don't need to be pruned if there's a star
 	_, selectStar := parentStars[table.Name()]
-	if unqualifiedStar {
-		selectStar = true
+	if selectStar || unqualifiedStar {
+		return n, transform.SameTree, nil
 	}
 
-	// Don't prune columns if they're needed by a virtual column
-	virtualColDeps := make(map[string]int)
-	if !selectStar { // if selectStar, we're adding all columns anyway
-		if vct, isVCT := n.WrappedTable().(*plan.VirtualColumnTable); isVCT {
-			for _, projection := range vct.Projections {
-				transform.InspectExpr(projection, func(e sql.Expression) bool {
-					if cd, isCD := e.(*sql.ColumnDefaultValue); isCD {
-						transform.InspectExpr(cd.Expr, func(e sql.Expression) bool {
-							if gf, ok := e.(*expression.GetField); ok {
-								virtualColDeps[gf.Name()]++
-							}
-							return false
-						})
-					}
-					return false
-				})
-			}
-		}
+	// pruning VirtualColumnTable underlying tables causes indexing errors when VirtualColumnTable.Projections (which are sql.Expression)
+	// are evaluated
+	if _, isVCT := n.WrappedTable().(*plan.VirtualColumnTable); isVCT {
+		return n, transform.SameTree, nil
 	}
 
 	cols := make([]string, 0)
 	source := strings.ToLower(table.Name())
 	for _, col := range table.Schema() {
 		c := newTableCol(source, col.Name)
-		if selectStar || parentCols[c] > 0 || virtualColDeps[c.Name()] > 0 {
+		if parentCols[c] > 0 {
 			cols = append(cols, c.col)
 		}
 	}