add customer driver tests

f0ssel · f0ssel · commit dae69c715d75 · 2024-08-21T16:11:03.000Z
diff --git a/coderd/database/connector.go b/coderd/database/connector.go
@@ -1,19 +1,19 @@
 package database
 
 import (
-	"context"
 	"database/sql/driver"
 
 	"github.com/lib/pq"
 )
 
-// ConnectorCreator can create a driver.Connector.
+// ConnectorCreator is a driver.Driver that can create a driver.Connector.
 type ConnectorCreator interface {
+	driver.Driver
 	Connector(name string) (driver.Connector, error)
 }
 
-// DialerConnector can create a driver.Connector and set a pq.Dialer.
+// DialerConnector is a driver.Connector that can set a pq.Dialer.
 type DialerConnector interface {
-	Connect(context.Context) (driver.Conn, error)
+	driver.Connector
 	Dialer(dialer pq.Dialer)
 }
diff --git a/coderd/database/dbtestutil/driver.go b/coderd/database/dbtestutil/driver.go
@@ -0,0 +1,128 @@
+package dbtestutil
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"fmt"
+
+	"github.com/lib/pq"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/cryptorand"
+)
+
+var (
+	_ driver.Driver             = &Driver{}
+	_ database.ConnectorCreator = &Driver{}
+	_ database.DialerConnector  = &Connector{}
+)
+
+type Driver struct {
+	name        string
+	inner       driver.Driver
+	connections []driver.Conn
+	listeners   map[chan struct{}]chan struct{}
+}
+
+func Register() (*Driver, error) {
+	db, err := sql.Open("postgres", "")
+	if err != nil {
+		return nil, xerrors.Errorf("failed to open database: %w", err)
+	}
+
+	su, err := cryptorand.StringCharset(cryptorand.Alpha, 10)
+	if err != nil {
+		return nil, xerrors.Errorf("failed to generate random string: %w", err)
+	}
+
+	d := &Driver{
+		name:      fmt.Sprintf("postgres-test-%s", su),
+		inner:     db.Driver(),
+		listeners: make(map[chan struct{}]chan struct{}),
+	}
+
+	sql.Register(d.name, d)
+
+	return d, nil
+}
+
+func (d *Driver) Open(name string) (driver.Conn, error) {
+	conn, err := d.inner.Open(name)
+	if err != nil {
+		return nil, xerrors.Errorf("failed to open connection: %w", err)
+	}
+
+	d.AddConnection(conn)
+
+	return conn, nil
+}
+
+func (d *Driver) Connector(name string) (driver.Connector, error) {
+	return &Connector{
+		name:   name,
+		driver: d,
+	}, nil
+}
+
+func (d *Driver) Name() string {
+	return d.name
+}
+
+func (d *Driver) AddConnection(conn driver.Conn) {
+	d.connections = append(d.connections, conn)
+
+	for listener := range d.listeners {
+		d.listeners[listener] <- struct{}{}
+	}
+
+}
+
+func (d *Driver) WaitForConnection() {
+	ch := make(chan struct{})
+	d.listeners[ch] = ch
+	<-ch
+	delete(d.listeners, ch)
+}
+
+func (d *Driver) DropConnections() {
+	for _, conn := range d.connections {
+		_ = conn.Close()
+	}
+	d.connections = nil
+}
+
+type Connector struct {
+	name   string
+	driver *Driver
+	dialer pq.Dialer
+}
+
+func (c *Connector) Connect(_ context.Context) (driver.Conn, error) {
+	if c.dialer != nil {
+		conn, err := pq.DialOpen(c.dialer, c.name)
+		if err != nil {
+			return nil, xerrors.Errorf("failed to dial open connection: %w", err)
+		}
+
+		c.driver.AddConnection(conn)
+
+		return conn, nil
+	}
+
+	conn, err := c.driver.Open(c.name)
+	if err != nil {
+		return nil, xerrors.Errorf("failed to open connection: %w", err)
+	}
+
+	return conn, nil
+}
+
+func (c *Connector) Driver() driver.Driver {
+	return c.driver
+}
+
+func (c *Connector) Dialer(dialer pq.Dialer) {
+	c.dialer = dialer
+}
diff --git a/coderd/database/pubsub/pubsub.go b/coderd/database/pubsub/pubsub.go
@@ -455,7 +455,10 @@ func (p *PGPubsub) startListener(ctx context.Context, connectURL string) error {
 	}
 
 	// Set the dialer if the connector supports it.
-	if dc, ok := connector.(database.DialerConnector); ok {
+	dc, ok := connector.(database.DialerConnector)
+	if !ok {
+		p.logger.Critical(ctx, "connector does not support setting log dialer, database connection debug logs will be missing")
+	} else {
 		dc.Dialer(dialer)
 	}
 
diff --git a/coderd/database/pubsub/pubsub_test.go b/coderd/database/pubsub/pubsub_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
@@ -86,7 +87,7 @@ func TestPGPubsub_Metrics(t *testing.T) {
 	for i := range colossalData {
 		colossalData[i] = 'q'
 	}
-	unsub1, err := uut.Subscribe(event, func(ctx context.Context, message []byte) {
+	unsub1, err := uut.Subscribe(event, func(_ context.Context, message []byte) {
 		messageChannel <- message
 	})
 	require.NoError(t, err)
@@ -119,3 +120,79 @@ func TestPGPubsub_Metrics(t *testing.T) {
 			!testutil.PromCounterGathered(t, metrics, "coder_pubsub_latency_measure_errs_total")
 	}, testutil.WaitShort, testutil.IntervalFast)
 }
+
+func TestPGPubsubDriver(t *testing.T) {
+	t.Parallel()
+	if !dbtestutil.WillUsePostgres() {
+		t.Skip("test only with postgres")
+	}
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	logger := slogtest.Make(t, &slogtest.Options{
+		IgnoreErrors: true,
+	}).Leveled(slog.LevelDebug)
+
+	connectionURL, closePg, err := dbtestutil.Open()
+	require.NoError(t, err)
+	defer closePg()
+
+	// wrap the pg driver with one we can control
+	d, err := dbtestutil.Register()
+	require.NoError(t, err)
+
+	db, err := sql.Open(d.Name(), connectionURL)
+	require.NoError(t, err)
+	defer db.Close()
+
+	ps, err := pubsub.New(ctx, logger, db, connectionURL)
+	require.NoError(t, err)
+	defer ps.Close()
+
+	// test that we can publish and subscribe
+	gotChan := make(chan struct{})
+	defer close(gotChan)
+	subCancel, err := ps.Subscribe("test", func(_ context.Context, _ []byte) {
+		gotChan <- struct{}{}
+	})
+	require.NoError(t, err)
+	defer subCancel()
+
+	err = ps.Publish("test", []byte("hello"))
+	require.NoError(t, err)
+
+	select {
+	case <-gotChan:
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for message")
+	}
+
+	reconnectChan := make(chan struct{})
+	go func() {
+		d.WaitForConnection()
+		// wait a bit to make sure the pubsub has reestablished it's connection
+		// if we don't wait, the publish may be dropped because the pubsub hasn't initialized yet.
+		time.Sleep(1 * time.Second)
+		reconnectChan <- struct{}{}
+	}()
+
+	// drop the underlying connection being used by the pubsub
+	// the pq.Listener should reconnect and repopulate it's listeners
+	// so old subscriptions should still work
+	d.DropConnections()
+
+	select {
+	case <-reconnectChan:
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for reconnect")
+	}
+
+	// ensure our old subscription still fires
+	err = ps.Publish("test", []byte("hello-again"))
+	require.NoError(t, err)
+
+	select {
+	case <-gotChan:
+	case <-ctx.Done():
+		t.Fatal("timeout waiting for message after reconnect")
+	}
+}

Original file line number	Diff line number	Diff line change
`@@ -455,7 +455,10 @@ func (p *PGPubsub) startListener(ctx context.Context, connectURL string) error {`
`455`	`455`	`}`
`456`	`456`
`457`	`457`	`// Set the dialer if the connector supports it.`
`458`		`- if dc, ok := connector.(database.DialerConnector); ok {`
	`458`	`+ dc, ok := connector.(database.DialerConnector)`
	`459`	`+ if !ok {`
	`460`	`+ p.logger.Critical(ctx, "connector does not support setting log dialer, database connection debug logs will be missing")`
	`461`	`+ } else {`
`459`	`462`	`dc.Dialer(dialer)`
`460`	`463`	`}`
`461`	`464`