storage/aws: accept a caller-owned MySQL DB

storage/aws/antispam/aws.go

@@ -35,7 +35,7 @@ import (
 	"github.com/transparency-dev/tessera/internal/otel"
 	"go.opentelemetry.io/otel/trace"
 
-	_ "github.com/go-sql-driver/mysql"
+	"github.com/go-sql-driver/mysql"
 )
 
 const (
@@ -71,6 +71,11 @@ type AntispamOpts struct {
 	// the antispam decorator will return a wrapped tessera.ErrPushback for every Add request.
 	PushbackThreshold uint
 
+	// DB will be used to interact with MySQL. If unset, Tessera will open one using the
+	// dsn parameter to NewAntispam, MaxOpenConns and MaxIdleConns. If set, it must be
+	// connected to the database named by dsn, and it will never be closed by Tessera.
+	DB *sql.DB
+
 	PushbackMaxOutstanding uint64
 	MaxOpenConns           int
 	MaxIdleConns           int
@@ -107,6 +112,13 @@ func NewAntispam(ctx context.Context, dsn string, opts AntispamOpts) (*AntispamS
 		opts.PushbackThreshold = DefaultPushbackThreshold
 	}
 
+	if opts.DB != nil {
+		if err := checkDatabase(ctx, opts.DB, dsn); err != nil {
+			return nil, err
+		}
+		return newAntispam(ctx, opts.DB, opts)
+	}
+
 	dbPool, err := sql.Open("mysql", dsn)
 	if err != nil {
 		return nil, fmt.Errorf("failed to connect to MySQL db: %v", err)
@@ -118,7 +130,12 @@ func NewAntispam(ctx context.Context, dsn string, opts AntispamOpts) (*AntispamS
 	if opts.MaxIdleConns >= 0 {
 		dbPool.SetMaxIdleConns(opts.MaxIdleConns)
 	}
+	return newAntispam(ctx, dbPool, opts)
+}
 
+// newAntispam returns an antispam driver which uses the provided database handle
+// for its MySQL connections.
+func newAntispam(ctx context.Context, dbPool *sql.DB, opts AntispamOpts) (*AntispamStorage, error) {
 	if err := dbPool.Ping(); err != nil {
 		return nil, fmt.Errorf("failed to ping MySQL db: %v", err)
 	}
@@ -137,6 +154,24 @@ func NewAntispam(ctx context.Context, dsn string, opts AntispamOpts) (*AntispamS
 	return r, nil
 }
 
+// checkDatabase returns an error unless db is connected to the database named by dsn.
+// A pool connected to another database would otherwise keep a separate, and so
+// incomplete, record of the log's entries, and fail to detect duplicates.
+func checkDatabase(ctx context.Context, db *sql.DB, dsn string) error {
+	cfg, err := mysql.ParseDSN(dsn)
+	if err != nil {
+		return fmt.Errorf("failed to parse DSN: %v", err)
+	}
+	var got sql.NullString
+	if err := db.QueryRowContext(ctx, "SELECT DATABASE()").Scan(&got); err != nil {
+		return fmt.Errorf("failed to read the database of the provided DB: %v", err)
+	}
+	if got.String != cfg.DBName {
+		return fmt.Errorf("provided DB is connected to database %q, want %q", got.String, cfg.DBName)
+	}
+	return nil
+}
+
 func (s *AntispamStorage) initDB(ctx context.Context) error {
 	if _, err := s.dbPool.ExecContext(ctx,
 		`CREATE TABLE IF NOT EXISTS AntispamMeta (

storage/aws/antispam/aws_test.go

@@ -25,6 +25,7 @@ import (
 
 	"log/slog"
 
+	"github.com/go-sql-driver/mysql"
 	"github.com/transparency-dev/tessera"
 	"github.com/transparency-dev/tessera/api"
 	"github.com/transparency-dev/tessera/testonly"
@@ -42,16 +43,44 @@ func TestMain(m *testing.M) {
 }
 
 func TestAntispam(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		sharedDB bool
+	}{
+		{name: "own DB"},
+		{name: "shared DB", sharedDB: true},
+	} {
+		t.Run(test.name, func(t *testing.T) { testAntispam(t, test.sharedDB) })
+	}
+}
+
+// testAntispam drives the full append, follow and duplicate lookup cycle, through a DB
+// provided in AntispamOpts if sharedDB is set.
+func testAntispam(t *testing.T, sharedDB bool) {
 	ctx := t.Context()
 	if canSkipMySQLTest(t, ctx) {
 		slog.WarnContext(ctx, "MySQL not available, skipping", slog.String("name", t.Name()))
 		t.Skip("MySQL not available, skipping test")
 	}
 	mustDropTables(t, ctx)
-	as, err := NewAntispam(ctx, *mySQLURI, AntispamOpts{})
+	opts := AntispamOpts{}
+	if sharedDB {
+		db, err := sql.Open("mysql", *mySQLURI)
+		if err != nil {
+			t.Fatal(err)
+		}
+		// As documented on AntispamOpts.DB, the caller retains ownership of the DB's
+		// lifecycle. Cleanup runs after this function's defers have shut down the log.
+		t.Cleanup(func() { _ = db.Close() })
+		opts.DB = db
+	}
+	as, err := NewAntispam(ctx, *mySQLURI, opts)
 	if err != nil {
 		t.Fatal(err)
 	}
+	if sharedDB && as.dbPool != opts.DB {
+		t.Fatal("NewAntispam did not use the provided DB")
+	}
 	fl, shutdown := testonly.NewTestLog(t, tessera.NewAppendOptions().WithCheckpointInterval(time.Second))
 	defer func() {
 		if err := shutdown(t.Context()); err != nil {
@@ -107,6 +136,39 @@ func TestAntispam(t *testing.T) {
 	}
 }
 
+func TestAntispamSharedDBWrongDatabase(t *testing.T) {
+	ctx := t.Context()
+	if canSkipMySQLTest(t, ctx) {
+		slog.WarnContext(ctx, "MySQL not available, skipping", slog.String("name", t.Name()))
+		t.Skip("MySQL not available, skipping test")
+	}
+	// A second, empty database that the provided DB is connected to instead.
+	admin, err := sql.Open("mysql", *mySQLURI)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = admin.Close() }()
+	const other = "test_tessera_other"
+	if _, err := admin.ExecContext(ctx, "CREATE DATABASE IF NOT EXISTS "+other); err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _, _ = admin.ExecContext(context.Background(), "DROP DATABASE "+other) }()
+	cfg, err := mysql.ParseDSN(*mySQLURI)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cfg.DBName = other
+	db, err := sql.Open("mysql", cfg.FormatDSN())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = db.Close() }()
+
+	if _, err := NewAntispam(ctx, *mySQLURI, AntispamOpts{DB: db}); err == nil {
+		t.Error("NewAntispam accepted a DB connected to a different database, want error")
+	}
+}
+
 func TestAntispamPushbackRecovers(t *testing.T) {
 	ctx := t.Context()
 	if canSkipMySQLTest(t, ctx) {

storage/aws/aws.go

@@ -162,6 +162,10 @@ type Config struct {
 	MaxOpenConns int
 	// Maximum idle database connections in the connection pool.
 	MaxIdleConns int
+	// DB will be used to interact with MySQL. If unset, Tessera will open one using DSN,
+	// MaxOpenConns and MaxIdleConns. If set, those are ignored, and it will never be
+	// closed by Tessera.
+	DB *sql.DB
 
 	// HTTPClient will be used for other HTTP requests. If unset, Tessera will use the net/http DefaultClient.
 	HTTPClient *http.Client
@@ -196,7 +200,7 @@ func New(ctx context.Context, cfg Config) (tessera.Driver, error) {
 
 // Appender creates a new tessera.Appender lifecycle object.
 func (s *Storage) Appender(ctx context.Context, opts *tessera.AppendOptions) (*tessera.Appender, tessera.LogReader, error) {
-	seq, err := newMySQLSequencer(ctx, s.cfg.DSN, uint64(opts.PushbackMaxOutstanding()), s.cfg.MaxOpenConns, s.cfg.MaxIdleConns)
+	seq, err := s.sequencer(ctx, uint64(opts.PushbackMaxOutstanding()))
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create MySQL sequencer: %v", err)
 	}
@@ -553,7 +557,7 @@ func (s *Storage) MigrationWriter(ctx context.Context, opts *tessera.MigrationOp
 		},
 		entriesPath: opts.EntriesPath(),
 	}
-	seq, err := newMySQLSequencer(ctx, s.cfg.DSN, DefaultPushbackMaxOutstanding, s.cfg.MaxOpenConns, s.cfg.MaxIdleConns)
+	seq, err := s.sequencer(ctx, DefaultPushbackMaxOutstanding)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create MySQL sequencer: %v", err)
 	}
@@ -896,6 +900,15 @@ type mySQLSequencer struct {
 	maxOutstanding uint64
 }
 
+// sequencer returns a mySQLSequencer which uses the configured DB, or opens one
+// for the configured DSN.
+func (s *Storage) sequencer(ctx context.Context, maxOutstanding uint64) (*mySQLSequencer, error) {
+	if s.cfg.DB != nil {
+		return newMySQLSequencerWithDB(ctx, s.cfg.DB, maxOutstanding)
+	}
+	return newMySQLSequencer(ctx, s.cfg.DSN, maxOutstanding, s.cfg.MaxOpenConns, s.cfg.MaxIdleConns)
+}
+
 // newMySQLSequencer returns a new mysqlSequencer struct which uses the provided
 // DSN for its MySQL connection.
 func newMySQLSequencer(ctx context.Context, dsn string, maxOutstanding uint64, maxOpenConns, maxIdleConns int) (*mySQLSequencer, error) {
@@ -910,7 +923,12 @@ func newMySQLSequencer(ctx context.Context, dsn string, maxOutstanding uint64, m
 	if maxIdleConns >= 0 {
 		dbPool.SetMaxIdleConns(maxIdleConns)
 	}
+	return newMySQLSequencerWithDB(ctx, dbPool, maxOutstanding)
+}
 
+// newMySQLSequencerWithDB returns a new mysqlSequencer struct which uses the provided
+// database handle for its MySQL connections.
+func newMySQLSequencerWithDB(ctx context.Context, dbPool *sql.DB, maxOutstanding uint64) (*mySQLSequencer, error) {
 	if err := dbPool.Ping(); err != nil {
 		return nil, fmt.Errorf("failed to ping MySQL db: %v", err)
 	}

storage/aws/aws_test.go

@@ -210,6 +210,35 @@ func TestMySQLSequencerPushback(t *testing.T) {
 	}
 }
 
+// TestStorageUsesProvidedDB checks that Storage sequences with the DB set in its
+// Config, rather than opening one for its DSN.
+func TestStorageUsesProvidedDB(t *testing.T) {
+	ctx := context.Background()
+	if canSkipMySQLTest(t, ctx) {
+		slog.WarnContext(ctx, "MySQL not available, skipping", slog.String("name", t.Name()))
+		t.Skip("MySQL not available, skipping test")
+	}
+	mustDropTables(t, ctx)
+	db, err := sql.Open("mysql", *mySQLURI)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// As documented on Config.DB, the caller retains ownership of the DB's lifecycle.
+	defer func() { _ = db.Close() }()
+
+	s := &Storage{cfg: Config{DB: db, DSN: "nobody:nothing@tcp(127.0.0.1:1)/nowhere"}}
+	seq, err := s.sequencer(ctx, 1000)
+	if err != nil {
+		t.Fatalf("sequencer: %v", err)
+	}
+	if seq.dbPool != db {
+		t.Fatal("the sequencer does not use the provided DB")
+	}
+	if err := seq.assignEntries(ctx, []*tessera.Entry{tessera.NewEntry([]byte("entry"))}); err != nil {
+		t.Fatalf("assignEntries: %v", err)
+	}
+}
+
 func TestMySQLSequencerRoundTrip(t *testing.T) {
 	ctx := context.Background()
 	if canSkipMySQLTest(t, ctx) {