6 files changed, 17 insertions, 13 deletions
diff --git a/keyserver/keyserver.go b/keyserver/keyserver.go
index c557dfba..007a48a5 100644
--- a/keyserver/keyserver.go
+++ b/keyserver/keyserver.go
@@ -41,7 +41,7 @@ func NewInternalAPI(
 ) api.KeyInternalAPI {
 	js, _ := jetstream.Prepare(base.ProcessContext, &cfg.Matrix.JetStream)
 
-	db, err := storage.NewDatabase(&cfg.Database)
+	db, err := storage.NewDatabase(base, &cfg.Database)
 	if err != nil {
 		logrus.WithError(err).Panicf("failed to connect to key server database")
 	}
diff --git a/keyserver/storage/postgres/storage.go b/keyserver/storage/postgres/storage.go
index d4c7e2cc..b8f70acf 100644
--- a/keyserver/storage/postgres/storage.go
+++ b/keyserver/storage/postgres/storage.go
@@ -18,13 +18,14 @@ import (
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/keyserver/storage/postgres/deltas"
 	"github.com/matrix-org/dendrite/keyserver/storage/shared"
+	"github.com/matrix-org/dendrite/setup/base"
 	"github.com/matrix-org/dendrite/setup/config"
 )
 
 // NewDatabase creates a new sync server database
-func NewDatabase(dbProperties *config.DatabaseOptions) (*shared.Database, error) {
+func NewDatabase(base *base.BaseDendrite, dbProperties *config.DatabaseOptions) (*shared.Database, error) {
 	var err error
-	db, err := sqlutil.Open(dbProperties)
+	db, writer, err := base.DatabaseConnection(dbProperties, sqlutil.NewDummyWriter())
 	if err != nil {
 		return nil, err
 	}
@@ -63,7 +64,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*shared.Database, error)
 	}
 	d := &shared.Database{
 		DB:                    db,
-		Writer:                sqlutil.NewDummyWriter(),
+		Writer:                writer,
 		OneTimeKeysTable:      otk,
 		DeviceKeysTable:       dk,
 		KeyChangesTable:       kc,
diff --git a/keyserver/storage/sqlite3/storage.go b/keyserver/storage/sqlite3/storage.go
index 84d4cdf5..aeea9eac 100644
--- a/keyserver/storage/sqlite3/storage.go
+++ b/keyserver/storage/sqlite3/storage.go
@@ -18,11 +18,12 @@ import (
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/keyserver/storage/shared"
 	"github.com/matrix-org/dendrite/keyserver/storage/sqlite3/deltas"
+	"github.com/matrix-org/dendrite/setup/base"
 	"github.com/matrix-org/dendrite/setup/config"
 )
 
-func NewDatabase(dbProperties *config.DatabaseOptions) (*shared.Database, error) {
-	db, err := sqlutil.Open(dbProperties)
+func NewDatabase(base *base.BaseDendrite, dbProperties *config.DatabaseOptions) (*shared.Database, error) {
+	db, writer, err := base.DatabaseConnection(dbProperties, sqlutil.NewExclusiveWriter())
 	if err != nil {
 		return nil, err
 	}
@@ -62,7 +63,7 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*shared.Database, error)
 	}
 	d := &shared.Database{
 		DB:                    db,
-		Writer:                sqlutil.NewExclusiveWriter(),
+		Writer:                writer,
 		OneTimeKeysTable:      otk,
 		DeviceKeysTable:       dk,
 		KeyChangesTable:       kc,
diff --git a/keyserver/storage/storage.go b/keyserver/storage/storage.go
index 742e8463..ab6a3540 100644
--- a/keyserver/storage/storage.go
+++ b/keyserver/storage/storage.go
@@ -22,17 +22,18 @@ import (
 
 	"github.com/matrix-org/dendrite/keyserver/storage/postgres"
 	"github.com/matrix-org/dendrite/keyserver/storage/sqlite3"
+	"github.com/matrix-org/dendrite/setup/base"
 	"github.com/matrix-org/dendrite/setup/config"
 )
 
 // NewDatabase opens a new Postgres or Sqlite database (based on dataSourceName scheme)
 // and sets postgres connection parameters
-func NewDatabase(dbProperties *config.DatabaseOptions) (Database, error) {
+func NewDatabase(base *base.BaseDendrite, dbProperties *config.DatabaseOptions) (Database, error) {
 	switch {
 	case dbProperties.ConnectionString.IsSQLite():
-		return sqlite3.NewDatabase(dbProperties)
+		return sqlite3.NewDatabase(base, dbProperties)
 	case dbProperties.ConnectionString.IsPostgres():
-		return postgres.NewDatabase(dbProperties)
+		return postgres.NewDatabase(base, dbProperties)
 	default:
 		return nil, fmt.Errorf("unexpected database type")
 	}
diff --git a/keyserver/storage/storage_test.go b/keyserver/storage/storage_test.go
index 84d2098a..9940eac6 100644
--- a/keyserver/storage/storage_test.go
+++ b/keyserver/storage/storage_test.go
@@ -22,7 +22,7 @@ func MustCreateDatabase(t *testing.T) (Database, func()) {
 		log.Fatal(err)
 	}
 	t.Logf("Database %s", tmpfile.Name())
-	db, err := NewDatabase(&config.DatabaseOptions{
+	db, err := NewDatabase(nil, &config.DatabaseOptions{
 		ConnectionString: config.DataSource(fmt.Sprintf("file://%s", tmpfile.Name())),
 	})
 	if err != nil {
diff --git a/keyserver/storage/storage_wasm.go b/keyserver/storage/storage_wasm.go
index 8b31bfd0..75c9053e 100644
--- a/keyserver/storage/storage_wasm.go
+++ b/keyserver/storage/storage_wasm.go
@@ -18,13 +18,14 @@ import (
 	"fmt"
 
 	"github.com/matrix-org/dendrite/keyserver/storage/sqlite3"
+	"github.com/matrix-org/dendrite/setup/base"
 	"github.com/matrix-org/dendrite/setup/config"
 )
 
-func NewDatabase(dbProperties *config.DatabaseOptions) (Database, error) {
+func NewDatabase(base *base.BaseDendrite, dbProperties *config.DatabaseOptions) (Database, error) {
 	switch {
 	case dbProperties.ConnectionString.IsSQLite():
-		return sqlite3.NewDatabase(dbProperties)
+		return sqlite3.NewDatabase(base, dbProperties)
 	case dbProperties.ConnectionString.IsPostgres():
 		return nil, fmt.Errorf("can't use Postgres implementation")
 	default: