Merge pull request #1452 from jamisonhyatt/fix-whitelist-table-blacklist-columns

Fix whitelist table blacklist columns

Author: stephenafamo

drivers/sqlboiler-mssql/driver/mssql.go

@@ -308,7 +308,9 @@ func (m *MSSQLDriver) Columns(schema, tableName string, whitelist, blacklist []s
 				args = append(args, w)
 			}
 		}
-	} else if len(blacklist) > 0 {
+	}
+	
+	if len(blacklist) > 0 {
 		cols := drivers.ColumnsFromList(blacklist, tableName)
 		if len(cols) > 0 {
 			query += fmt.Sprintf(" and c.column_name not in (%s)", strmangle.Placeholders(true, len(cols), 3, 1))

drivers/sqlboiler-mssql/driver/mssql_test.go

@@ -106,4 +106,36 @@ func TestDriver(t *testing.T) {
 	if bytes.Compare(want, got) != 0 {
 		t.Errorf("want:\n%s\ngot:\n%s\n", want, got)
 	}
+
+	t.Run("whitelist table, blacklist column", func(t *testing.T) {
+		p := &MSSQLDriver{}
+		config := drivers.Config{
+			"user":    envUsername,
+			"pass":    envPassword,
+			"dbname":  envDatabase,
+			"host":    envHostname,
+			"port":    envPort,
+			"sslmode": "disable",
+			"schema":  "dbo",
+			"whitelist": []string{"magic"},
+			"blacklist": []string{"magic.string_three"},
+		}
+		info, err := p.Assemble(config)
+		if err != nil {
+			t.Fatal(err)
+		}
+		found := false
+		for _, tbl := range info.Tables {
+			if tbl.Name == "magic" {
+				for _, col := range tbl.Columns {
+					if col.Name == "string_three" {
+						found = true
+					}
+				}
+			}
+		}
+		if found {
+			t.Errorf("blacklisted column 'string_three' should not be present in table 'magic'")
+		}
+	})
 }

drivers/sqlboiler-mysql/driver/mysql.go

@@ -332,7 +332,9 @@ func (m *MySQLDriver) Columns(schema, tableName string, whitelist, blacklist []s
 				args = append(args, w)
 			}
 		}
-	} else if len(blacklist) > 0 {
+	}
+	
+	if len(blacklist) > 0 {
 		cols := drivers.ColumnsFromList(blacklist, tableName)
 		if len(cols) > 0 {
 			query += fmt.Sprintf(" and c.column_name not in (%s)", strings.Repeat(",?", len(cols))[1:])

drivers/sqlboiler-mysql/driver/mysql_test.go

@@ -113,4 +113,32 @@ func TestDriver(t *testing.T) {
 			require.JSONEq(t, string(want), string(got))
 		})
 	}
+
+	t.Run("whitelist table, blacklist column", func(t *testing.T) {
+		p := &MySQLDriver{}
+		config := drivers.Config{
+			"user":    envUsername,
+			"pass":    envPassword,
+			"dbname":  envDatabase,
+			"host":    envHostname,
+			"port":    envPort,
+			"sslmode": "false",
+			"schema":  envDatabase,
+			"whitelist": []string{"magic"},
+			"blacklist": []string{"magic.string_three"},
+		}
+		info, err := p.Assemble(config)
+		require.NoError(t, err)
+		found := false
+		for _, tbl := range info.Tables {
+			if tbl.Name == "magic" {
+				for _, col := range tbl.Columns {
+					if col.Name == "string_three" {
+						found = true
+					}
+				}
+			}
+		}
+		require.False(t, found, "blacklisted column 'string_three' should not be present in table 'magic'")
+	})
 }

drivers/sqlboiler-psql/driver/psql.go

@@ -600,7 +600,9 @@ func (p *PostgresDriver) Columns(schema, tableName string, whitelist, blacklist
 				args = append(args, w)
 			}
 		}
-	} else if len(blacklist) > 0 {
+	}
+	
+	if len(blacklist) > 0 {
 		cols := drivers.ColumnsFromList(blacklist, tableName)
 		if len(cols) > 0 {
 			query += fmt.Sprintf(" where c.column_name not in (%s)", strmangle.Placeholders(true, len(cols), 3, 1))

drivers/sqlboiler-psql/driver/psql_test.go

@@ -114,4 +114,32 @@ func TestAssemble(t *testing.T) {
 			require.JSONEq(t, string(want), string(got))
 		})
 	}
+
+	t.Run("whitelist table, blacklist column", func(t *testing.T) {
+		p := PostgresDriver{}
+		config := drivers.Config{
+			"user":    envUsername,
+			"pass":    envPassword,
+			"dbname":  envDatabase,
+			"host":    envHostname,
+			"port":    envPort,
+			"sslmode": "disable",
+			"schema":  "public",
+			"whitelist": []string{"magic"},
+			"blacklist": []string{"magic.string_three"},
+		}
+		info, err := p.Assemble(config)
+		require.NoError(t, err)
+		found := false
+		for _, tbl := range info.Tables {
+			if tbl.Name == "magic" {
+				for _, col := range tbl.Columns {
+					if col.Name == "string_three" {
+						found = true
+					}
+				}
+			}
+		}
+		require.False(t, found, "blacklisted column 'string_three' should not be present in table 'magic'")
+	})
 }