Merge pull request #15049 from manadart/dqlite-bootstrap-options

[JUJU-2433] Allow Dqlite bootstrap to accept extra database mutations

database/bootstrap.go

@@ -5,6 +5,7 @@ package database
 
 import (
 	"context"
+	"database/sql"
 
 	"github.com/canonical/go-dqlite/app"
 	"github.com/juju/errors"
@@ -30,10 +31,13 @@ type bootstrapOptFactory interface {
 // BootstrapDqlite opens a new database for the controller, and runs the
 // DDL to create its schema.
 //
+// It accepts an optional list of functions to perform operations on the
+// controller database.
+//
 // At this point we know there are no peers and that we are the only user
 // of Dqlite, so we can eschew external address and clustering concerns.
 // Those will be handled by the db-accessor worker.
-func BootstrapDqlite(ctx context.Context, opt bootstrapOptFactory, logger Logger) error {
+func BootstrapDqlite(ctx context.Context, opt bootstrapOptFactory, logger Logger, ops ...func(db *sql.DB) error) error {
 	dir, err := opt.EnsureDataDir()
 	if err != nil {
 		return errors.Trace(err)
@@ -75,5 +79,11 @@ func BootstrapDqlite(ctx context.Context, opt bootstrapOptFactory, logger Logger
 		return errors.Annotate(err, "creating controller database schema")
 	}
 
+	for i, op := range ops {
+		if err := op(db); err != nil {
+			return errors.Annotatef(err, "running bootstrap operation at index %d", i)
+		}
+	}
+
 	return nil
 }

database/bootstrap_test.go

@@ -5,6 +5,8 @@ package database
 
 import (
 	"context"
+	"database/sql"
+	"errors"
 	"fmt"
 	"net"
 
@@ -24,35 +26,36 @@ var _ = gc.Suite(&bootstrapSuite{})
 func (s *bootstrapSuite) TestBootstrapSuccess(c *gc.C) {
 	opt := &testOptFactory{c: c}
 
-	c.Assert(BootstrapDqlite(context.TODO(), opt, stubLogger{}), jc.ErrorIsNil)
+	// check tests the variadic operation functionality
+	// and ensures that bootstrap applied the DDL.
+	check := func(db *sql.DB) error {
+		rows, err := db.Query("SELECT COUNT(*) FROM lease_type")
+		if err != nil {
+			return err
+		}
 
-	// Now use the same options to reopen the controller database
-	// and check that we can see bootstrap side effects.
-	dir, err := opt.EnsureDataDir()
-	c.Assert(err, jc.ErrorIsNil)
+		defer func() { _ = rows.Close() }()
 
-	addrOpt, err := opt.WithAddressOption()
-	c.Assert(err, jc.ErrorIsNil)
+		if !rows.Next() {
+			return errors.New("no rows in lease_type")
+		}
 
-	dqlite, err := app.New(dir, addrOpt)
-	c.Assert(err, jc.ErrorIsNil)
-	s.AddCleanup(func(*gc.C) { _ = dqlite.Close() })
+		var count int
+		err = rows.Scan(&count)
+		if err != nil {
+			return err
+		}
 
-	ctx := context.TODO()
-	c.Assert(dqlite.Ready(ctx), jc.ErrorIsNil)
+		if count != 3 {
+			return fmt.Errorf("expected 3 rows, got %d", count)
+		}
 
-	db, err := dqlite.Open(ctx, "controller")
-	c.Assert(err, jc.ErrorIsNil)
-	s.AddCleanup(func(*gc.C) { _ = db.Close() })
+		return nil
+	}
 
-	rows, err := db.Query("SELECT COUNT(*) FROM lease_type")
+	err := BootstrapDqlite(context.TODO(), opt, stubLogger{}, check)
 	c.Assert(err, jc.ErrorIsNil)
-	s.AddCleanup(func(*gc.C) { _ = rows.Close() })
 
-	var count int
-	c.Assert(rows.Next(), jc.IsTrue)
-	c.Assert(rows.Scan(&count), jc.ErrorIsNil)
-	c.Check(count, gc.Equals, 3)
 }
 
 type testOptFactory struct {