Update ConnectionManager to still allow component defined connections (#3154)

internal/sqlutil/connection_manager.go

@@ -17,16 +17,21 @@ package sqlutil
 import (
 	"database/sql"
 	"fmt"
+	"sync"
 
 	"github.com/matrix-org/dendrite/setup/config"
 	"github.com/matrix-org/dendrite/setup/process"
 )
 
 type Connections struct {
-	db             *sql.DB
-	writer         Writer
 	globalConfig        config.DatabaseOptions
 	processContext      *process.ProcessContext
+	existingConnections sync.Map
+}
+
+type con struct {
+	db     *sql.DB
+	writer Writer
 }
 
 func NewConnectionManager(processCtx *process.ProcessContext, globalConfig config.DatabaseOptions) *Connections {
@@ -38,9 +43,13 @@ func NewConnectionManager(processCtx *process.ProcessContext, globalConfig confi
 
 func (c *Connections) Connection(dbProperties *config.DatabaseOptions) (*sql.DB, Writer, error) {
 	var err error
+	// If no connectionString was provided, try the global one
 	if dbProperties.ConnectionString == "" {
-		// if no connectionString was provided, try the global one
 		dbProperties = &c.globalConfig
+		// If we still don't have a connection string, that's a problem
+		if dbProperties.ConnectionString == "" {
+			return nil, nil, fmt.Errorf("no database connections configured")
+		}
 	}
 
 	writer := NewDummyWriter()
@@ -48,13 +57,19 @@ func (c *Connections) Connection(dbProperties *config.DatabaseOptions) (*sql.DB,
 		writer = NewExclusiveWriter()
 	}
 
-	if dbProperties.ConnectionString != "" && c.db == nil {
+	existing, loaded := c.existingConnections.LoadOrStore(dbProperties.ConnectionString, &con{})
+	if loaded {
+		// We found an existing connection
+		ex := existing.(*con)
+		return ex.db, ex.writer, nil
+	}
+
 	// Open a new database connection using the supplied config.
-		c.db, err = Open(dbProperties, writer)
+	db, err := Open(dbProperties, writer)
 	if err != nil {
 		return nil, nil, err
 	}
-		c.writer = writer
+	c.existingConnections.Store(dbProperties.ConnectionString, &con{db: db, writer: writer})
 	go func() {
 		if c.processContext == nil {
 			return
@@ -63,15 +78,9 @@ func (c *Connections) Connection(dbProperties *config.DatabaseOptions) (*sql.DB,
 		// Dendrite to shut down to cleanly close the database connection.
 		c.processContext.ComponentStarted()
 		<-c.processContext.WaitForShutdown()
-			_ = c.db.Close()
+		_ = db.Close()
 		c.processContext.ComponentFinished()
 	}()
-		return c.db, c.writer, nil
+	return db, writer, nil
-	}
+
-	if c.db != nil && c.writer != nil {
-		// Ignore the supplied config and return the global pool and
-		// writer.
-		return c.db, c.writer, nil
-	}
-	return nil, nil, fmt.Errorf("no database connections configured")
 }

internal/sqlutil/connection_manager_test.go

@@ -48,6 +48,22 @@ func TestConnectionManager(t *testing.T) {
 			if !reflect.DeepEqual(writer, writer2) {
 				t.Fatalf("expected database writer to be reused")
 			}
+
+			// This test does not work with Postgres, because we can't just simply append
+			// "x" or replace the database to use.
+			if dbType == test.DBTypePostgres {
+				return
+			}
+
+			// Test different connection string
+			dbProps = &config.DatabaseOptions{ConnectionString: config.DataSource(conStr + "x")}
+			db3, _, err := cm.Connection(dbProps)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if reflect.DeepEqual(db, db3) {
+				t.Fatalf("expected different database connection")
+			}
 		})
 	})
 
@@ -115,4 +131,10 @@ func TestConnectionManager(t *testing.T) {
 	if err == nil {
 		t.Fatal("expected an error but got none")
 	}
+
+	// empty connection string is not allowed
+	_, _, err = cm2.Connection(&config.DatabaseOptions{})
+	if err == nil {
+		t.Fatal("expected an error but got none")
+	}
 }