Add whitelist feature

bdb/drivers/postgres.go

@@ -9,7 +9,6 @@ import (
 	_ "github.com/lib/pq"
 	"github.com/pkg/errors"
 	"github.com/vattle/sqlboiler/bdb"
-	"github.com/vattle/sqlboiler/strmangle"
 )
 
 // PostgresDriver holds the database connection string and a handle
@@ -82,18 +81,15 @@ func (p *PostgresDriver) UseLastInsertID() bool {
 
 // TableNames connects to the postgres database and
 // retrieves all table names from the information_schema where the
-// table schema is public. It excludes common migration tool tables
-// such as gorp_migrations
-func (p *PostgresDriver) TableNames(exclude []string) ([]string, error) {
+// table schema is public. It uses a whitelist and exclude list.
+func (p *PostgresDriver) TableNames(whitelist, exclude []string) ([]string, error) {
 	var names []string
 
 	query := `select table_name from information_schema.tables where table_schema = 'public'`
-	if len(exclude) > 0 {
-		quoteStr := func(x string) string {
-			return `'` + x + `'`
-		}
-		exclude = strmangle.StringMap(quoteStr, exclude)
-		query = query + fmt.Sprintf("and table_name not in (%s);", strings.Join(exclude, ","))
+	if len(whitelist) > 0 {
+		query = query + fmt.Sprintf("and table_name in ('%s');", strings.Join(whitelist, "','"))
+	} else if len(exclude) > 0 {
+		query = query + fmt.Sprintf("and table_name not in ('%s');", strings.Join(exclude, "','"))
 	}
 
 	rows, err := p.dbConn.Query(query)