diff --git a/go/libraries/doltcore/sqle/database.go b/go/libraries/doltcore/sqle/database.go index 54be8c5f5a..730934c84a 100644 --- a/go/libraries/doltcore/sqle/database.go +++ b/go/libraries/doltcore/sqle/database.go @@ -37,19 +37,33 @@ const ( // Database implements sql.Database for a dolt DB. type Database struct { - sql.Database - name string - root *doltdb.RootValue - dEnv *env.DoltEnv - batchMode batchMode + name string + root *doltdb.RootValue + dEnv *env.DoltEnv + batchMode batchMode + outstandingTables map[string]*DoltTable } // NewDatabase returns a new dolt database to use in queries. func NewDatabase(name string, root *doltdb.RootValue, dEnv *env.DoltEnv) *Database { return &Database{ - name: name, - root: root, - dEnv: dEnv, + name: name, + root: root, + dEnv: dEnv, + batchMode: single, + outstandingTables: make(map[string]*DoltTable), + } +} + +// NewBatchedDatabase returns a new dolt database executing in batch insert mode. Integrators must call Flush() to +// commit any outstanding edits. +func NewBatchedDatabase(name string, root *doltdb.RootValue, dEnv *env.DoltEnv) *Database { + return &Database{ + name: name, + root: root, + dEnv: dEnv, + batchMode: batched, + outstandingTables: make(map[string]*DoltTable), } } @@ -101,7 +115,12 @@ func (db *Database) GetTableInsensitive(ctx context.Context, tblName string) (sq return nil, false, err } - return &DoltTable{name: tblName, table: tbl, sch: sch, db: db}, true, nil + table := &DoltTable{name: tblName, table: tbl, sch: sch, db: db} + if db.batchMode == batched { + db.outstandingTables[tblName] = table + } + + return table, true, nil } func (db *Database) GetTableNames(ctx context.Context) ([]string, error) { @@ -183,3 +202,13 @@ func (db *Database) CreateTable(ctx *sql.Context, tableName string, schema sql.S return nil } + +// Flushes the current batch of outstanding changes and returns any errors. +func (db *Database) Flush(ctx context.Context) error { + for _, table := range db.outstandingTables { + if err := table.flushBatchedEdits(ctx); err != nil { + return err + } + } + return nil +} \ No newline at end of file