Merge pull request volatiletech#1143 from pavel-krush/master

Fetch column and table info in parallel

Author: stephenafamo

drivers/interface.go

@@ -4,6 +4,7 @@ package drivers
 
 import (
 	"sort"
+	"sync"
 
 	"github.com/friendsofgo/errors"
 	"github.com/volatiletech/sqlboiler/v4/importers"
@@ -17,13 +18,17 @@ const (
 	ConfigSchema         = "schema"
 	ConfigAddEnumTypes   = "add-enum-types"
 	ConfigEnumNullPrefix = "enum-null-prefix"
+	ConfigConcurrency    = "concurrency"
 
 	ConfigUser    = "user"
 	ConfigPass    = "pass"
 	ConfigHost    = "host"
 	ConfigPort    = "port"
 	ConfigDBName  = "dbname"
 	ConfigSSLMode = "sslmode"
+
+	// DefaultConcurrency defines the default amount of threads to use when loading tables info
+	DefaultConcurrency = 10
 )
 
 // Interface abstracts either a side-effect imported driver or a binary
@@ -102,77 +107,122 @@ type TableColumnTypeTranslator interface {
 // Tables returns the metadata for all tables, minus the tables
 // specified in the blacklist.
 func Tables(c Constructor, schema string, whitelist, blacklist []string) ([]Table, error) {
+	return TablesConcurrently(c, schema, whitelist, blacklist, 1)
+}
+
+// TablesConcurrently is a concurrent version of Tables. It returns the
+// metadata for all tables, minus the tables specified in the blacklist.
+func TablesConcurrently(c Constructor, schema string, whitelist, blacklist []string, concurrency int) ([]Table, error) {
 	var err error
+	var ret []Table
 
-	names, err := c.TableNames(schema, whitelist, blacklist)
+	ret, err = tables(c, schema, whitelist, blacklist, concurrency)
 	if err != nil {
-		return nil, errors.Wrap(err, "unable to get table names")
+		return nil, errors.Wrap(err, "unable to load tables")
 	}
 
-	sort.Strings(names)
-
-	var tables []Table
-	for _, name := range names {
-		t := Table{
-			Name: name,
+	if vc, ok := c.(ViewConstructor); ok {
+		v, err := views(vc, schema, whitelist, blacklist, concurrency)
+		if err != nil {
+			return nil, errors.Wrap(err, "unable to load views")
 		}
+		ret = append(ret, v...)
+	}
 
-		if t.Columns, err = c.Columns(schema, name, whitelist, blacklist); err != nil {
-			return nil, errors.Wrapf(err, "unable to fetch table column info (%s)", name)
-		}
+	return ret, nil
+}
 
-		tr, ok := c.(TableColumnTypeTranslator)
-		if ok {
-			for i, col := range t.Columns {
-				t.Columns[i] = tr.TranslateTableColumnType(col, name)
-			}
-		} else {
-			for i, col := range t.Columns {
-				t.Columns[i] = c.TranslateColumnType(col)
-			}
-		}
+func tables(c Constructor, schema string, whitelist, blacklist []string, concurrency int) ([]Table, error) {
+	var err error
 
-		if t.PKey, err = c.PrimaryKeyInfo(schema, name); err != nil {
-			return nil, errors.Wrapf(err, "unable to fetch table pkey info (%s)", name)
-		}
+	names, err := c.TableNames(schema, whitelist, blacklist)
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to get table names")
+	}
 
-		if t.FKeys, err = c.ForeignKeyInfo(schema, name); err != nil {
-			return nil, errors.Wrapf(err, "unable to fetch table fkey info (%s)", name)
-		}
+	sort.Strings(names)
 
-		filterPrimaryKey(&t, whitelist, blacklist)
-		filterForeignKeys(&t, whitelist, blacklist)
+	ret := make([]Table, len(names))
+
+	limiter := newConcurrencyLimiter(concurrency)
+	wg := sync.WaitGroup{}
+	errs := make(chan error, len(names))
+	for i, name := range names {
+		wg.Add(1)
+		limiter.get()
+		go func(i int, name string) {
+			defer wg.Done()
+			defer limiter.put()
+			t, err := table(c, schema, name, whitelist, blacklist)
+			if err != nil {
+				errs <- err
+				return
+			}
+			ret[i] = t
+		}(i, name)
+	}
 
-		setIsJoinTable(&t)
+	wg.Wait()
 
-		tables = append(tables, t)
+	// return first error occurred if any
+	if len(errs) > 0 {
+		return nil, <-errs
 	}
 
 	// Relationships have a dependency on foreign key nullability.
-	for i := range tables {
-		tbl := &tables[i]
-		setForeignKeyConstraints(tbl, tables)
+	for i := range ret {
+		tbl := &ret[i]
+		setForeignKeyConstraints(tbl, ret)
 	}
-	for i := range tables {
-		tbl := &tables[i]
-		setRelationships(tbl, tables)
+	for i := range ret {
+		tbl := &ret[i]
+		setRelationships(tbl, ret)
 	}
 
-	if vc, ok := c.(ViewConstructor); ok {
-		viewTables, err := views(vc, schema, whitelist, blacklist)
-		if err != nil {
-			return nil, errors.Wrap(err, "unable to load views")
+	return ret, nil
+}
+
+// table returns columns info for a given table
+func table(c Constructor, schema string, name string, whitelist, blacklist []string) (Table, error) {
+	var err error
+	t := &Table{
+		Name: name,
+	}
+
+	if t.Columns, err = c.Columns(schema, name, whitelist, blacklist); err != nil {
+		return Table{}, errors.Wrapf(err, "unable to fetch table column info (%s)", name)
+	}
+
+	tr, ok := c.(TableColumnTypeTranslator)
+	if ok {
+		for i, col := range t.Columns {
+			t.Columns[i] = tr.TranslateTableColumnType(col, name)
 		}
+	} else {
+		for i, col := range t.Columns {
+			t.Columns[i] = c.TranslateColumnType(col)
+		}
+	}
 
-		tables = append(tables, viewTables...)
+	if t.PKey, err = c.PrimaryKeyInfo(schema, name); err != nil {
+		return Table{}, errors.Wrapf(err, "unable to fetch table pkey info (%s)", name)
 	}
 
-	return tables, nil
+	if t.FKeys, err = c.ForeignKeyInfo(schema, name); err != nil {
+		return Table{}, errors.Wrapf(err, "unable to fetch table fkey info (%s)", name)
+	}
+
+	filterPrimaryKey(t, whitelist, blacklist)
+	filterForeignKeys(t, whitelist, blacklist)
+
+	setIsJoinTable(t)
+
+	return *t, nil
 }
 
 // views returns the metadata for all views, minus the views
 // specified in the blacklist.
-func views(c ViewConstructor, schema string, whitelist, blacklist []string) ([]Table, error) {
+func views(c ViewConstructor, schema string, whitelist, blacklist []string, concurrency int) ([]Table, error) {
 	var err error
 
 	names, err := c.ViewNames(schema, whitelist, blacklist)
@@ -182,44 +232,71 @@ func views(c ViewConstructor, schema string, whitelist, blacklist []string) ([]T
 
 	sort.Strings(names)
 
-	var views []Table
-	for _, name := range names {
-		t := Table{
-			IsView: true,
-			Name:   name,
-		}
+	ret := make([]Table, len(names))
+
+	limiter := newConcurrencyLimiter(concurrency)
+	wg := sync.WaitGroup{}
+	errs := make(chan error, len(names))
+	for i, name := range names {
+		wg.Add(1)
+		limiter.get()
+		go func(i int, name string) {
+			defer wg.Done()
+			defer limiter.put()
+			t, err := view(c, schema, name, whitelist, blacklist)
+			if err != nil {
+				errs <- err
+				return
+			}
+			ret[i] = t
+		}(i, name)
+	}
 
-		if t.ViewCapabilities, err = c.ViewCapabilities(schema, name); err != nil {
-			return nil, errors.Wrapf(err, "unable to fetch view capabilities info (%s)", name)
-		}
+	wg.Wait()
 
-		if t.Columns, err = c.ViewColumns(schema, name, whitelist, blacklist); err != nil {
-			return nil, errors.Wrapf(err, "unable to fetch view column info (%s)", name)
-		}
+	// return first error occurred if any
+	if len(errs) > 0 {
+		return nil, <-errs
+	}
 
-		tr, ok := c.(TableColumnTypeTranslator)
-		if ok {
-			for i, col := range t.Columns {
-				t.Columns[i] = tr.TranslateTableColumnType(col, name)
-			}
-		} else {
-			for i, col := range t.Columns {
-				t.Columns[i] = c.TranslateColumnType(col)
-			}
-		}
+	return ret, nil
+}
 
-		views = append(views, t)
+// view returns columns info for a given view
+func view(c ViewConstructor, schema string, name string, whitelist, blacklist []string) (Table, error) {
+	var err error
+	t := Table{
+		IsView: true,
+		Name:   name,
 	}
 
-	return views, nil
+	if t.ViewCapabilities, err = c.ViewCapabilities(schema, name); err != nil {
+		return Table{}, errors.Wrapf(err, "unable to fetch view capabilities info (%s)", name)
+	}
+
+	if t.Columns, err = c.ViewColumns(schema, name, whitelist, blacklist); err != nil {
+		return Table{}, errors.Wrapf(err, "unable to fetch view column info (%s)", name)
+	}
+
+	tr, ok := c.(TableColumnTypeTranslator)
+	if ok {
+		for i, col := range t.Columns {
+			t.Columns[i] = tr.TranslateTableColumnType(col, name)
+		}
+	} else {
+		for i, col := range t.Columns {
+			t.Columns[i] = c.TranslateColumnType(col)
+		}
+	}
+
+	return t, nil
 }
 
 func knownColumn(table string, column string, whitelist, blacklist []string) bool {
 	return (len(whitelist) == 0 ||
 		strmangle.SetInclude(table, whitelist) ||
 		strmangle.SetInclude(table+"."+column, whitelist) ||
 		strmangle.SetInclude("*."+column, whitelist)) &&
-
 		(len(blacklist) == 0 || (!strmangle.SetInclude(table, blacklist) &&
 			!strmangle.SetInclude(table+"."+column, blacklist) &&
 			!strmangle.SetInclude("*."+column, blacklist)))
@@ -294,3 +371,23 @@ func setRelationships(t *Table, tables []Table) {
 	t.ToOneRelationships = toOneRelationships(*t, tables)
 	t.ToManyRelationships = toManyRelationships(*t, tables)
 }
+
+// concurrencyCounter is a helper structure that can limit amount of concurrently processed requests
+type concurrencyLimiter chan struct{}
+
+func newConcurrencyLimiter(capacity int) concurrencyLimiter {
+	ret := make(concurrencyLimiter, capacity)
+	for i := 0; i < capacity; i++ {
+		ret <- struct{}{}
+	}
+
+	return ret
+}
+
+func (c concurrencyLimiter) get() {
+	<-c
+}
+
+func (c concurrencyLimiter) put() {
+	c <- struct{}{}
+}

drivers/interface_test.go

@@ -115,7 +115,7 @@ func (m testMockDriver) UseIndexPlaceholders() bool {
 func TestTables(t *testing.T) {
 	t.Parallel()
 
-	tables, err := Tables(testMockDriver{}, "public", nil, nil)
+	tables, err := TablesConcurrently(testMockDriver{}, "public", nil, nil, 1)
 	if err != nil {
 		t.Error(err)
 	}

drivers/mocks/mock.go

@@ -66,7 +66,7 @@ func (m *MockDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo, er
 	whitelist, _ := config.StringSlice(drivers.ConfigWhitelist)
 	blacklist, _ := config.StringSlice(drivers.ConfigBlacklist)
 
-	dbinfo.Tables, err = drivers.Tables(m, schema, whitelist, blacklist)
+	dbinfo.Tables, err = drivers.TablesConcurrently(m, schema, whitelist, blacklist, 1)
 	if err != nil {
 		return nil, err
 	}

drivers/sqlboiler-mssql/driver/mssql.go

@@ -81,6 +81,7 @@ func (m *MSSQLDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo, e
 	schema := config.DefaultString(drivers.ConfigSchema, "dbo")
 	whitelist, _ := config.StringSlice(drivers.ConfigWhitelist)
 	blacklist, _ := config.StringSlice(drivers.ConfigBlacklist)
+	concurrency := config.DefaultInt(drivers.ConfigConcurrency, drivers.DefaultConcurrency)
 
 	m.connStr = MSSQLBuildQueryString(user, pass, dbname, host, port, sslmode)
 	m.conn, err = sql.Open("mssql", m.connStr)
@@ -110,7 +111,7 @@ func (m *MSSQLDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo, e
 			UseCaseWhenExistsClause: true,
 		},
 	}
-	dbinfo.Tables, err = drivers.Tables(m, schema, whitelist, blacklist)
+	dbinfo.Tables, err = drivers.TablesConcurrently(m, schema, whitelist, blacklist, concurrency)
 	if err != nil {
 		return nil, err
 	}

drivers/sqlboiler-mysql/driver/mysql.go

@@ -83,6 +83,7 @@ func (m *MySQLDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo, e
 	schema := dbname
 	whitelist, _ := config.StringSlice(drivers.ConfigWhitelist)
 	blacklist, _ := config.StringSlice(drivers.ConfigBlacklist)
+	concurrency := config.DefaultInt(drivers.ConfigConcurrency, drivers.DefaultConcurrency)
 
 	tinyIntAsIntIntf, ok := config["tinyint_as_int"]
 	if ok {
@@ -116,7 +117,7 @@ func (m *MySQLDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo, e
 		},
 	}
 
-	dbinfo.Tables, err = drivers.Tables(m, schema, whitelist, blacklist)
+	dbinfo.Tables, err = drivers.TablesConcurrently(m, schema, whitelist, blacklist, concurrency)
 	if err != nil {
 		return nil, err
 	}

drivers/sqlboiler-psql/driver/psql.go

@@ -96,6 +96,7 @@ func (p *PostgresDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo
 	schema := config.DefaultString(drivers.ConfigSchema, "public")
 	whitelist, _ := config.StringSlice(drivers.ConfigWhitelist)
 	blacklist, _ := config.StringSlice(drivers.ConfigBlacklist)
+	concurrency := config.DefaultInt(drivers.ConfigConcurrency, drivers.DefaultConcurrency)
 
 	useSchema := schema != "public"
 
@@ -130,7 +131,7 @@ func (p *PostgresDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo
 			UseDefaultKeyword:    true,
 		},
 	}
-	dbinfo.Tables, err = drivers.Tables(p, schema, whitelist, blacklist)
+	dbinfo.Tables, err = drivers.TablesConcurrently(p, schema, whitelist, blacklist, concurrency)
 	if err != nil {
 		return nil, err
 	}

drivers/sqlboiler-sqlite3/driver/sqlite3.go

@@ -70,6 +70,7 @@ func (s SQLiteDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo, e
 	dbname := config.MustString(drivers.ConfigDBName)
 	whitelist, _ := config.StringSlice(drivers.ConfigWhitelist)
 	blacklist, _ := config.StringSlice(drivers.ConfigBlacklist)
+	concurrency := config.DefaultInt(drivers.ConfigConcurrency, drivers.DefaultConcurrency)
 
 	s.connStr = SQLiteBuildQueryString(dbname)
 	s.dbConn, err = sql.Open("sqlite", s.connStr)
@@ -95,7 +96,7 @@ func (s SQLiteDriver) Assemble(config drivers.Config) (dbinfo *drivers.DBInfo, e
 		},
 	}
 
-	dbinfo.Tables, err = drivers.Tables(s, "", whitelist, blacklist)
+	dbinfo.Tables, err = drivers.TablesConcurrently(s, "", whitelist, blacklist, concurrency)
 	if err != nil {
 		return nil, err
 	}