Implement interfaces for keydb storage

common/basecomponent/base.go

@@ -138,7 +138,7 @@ func (b *BaseDendrite) CreateAccountsDB() *accounts.Database {
 
 // CreateKeyDB creates a new instance of the key database. Should only be called
 // once per component.
-func (b *BaseDendrite) CreateKeyDB() *keydb.Database {
+func (b *BaseDendrite) CreateKeyDB() keydb.Database {
 	db, err := keydb.NewDatabase(string(b.Cfg.Database.ServerKey))
 	if err != nil {
 		logrus.WithError(err).Panicf("failed to connect to keys db")

common/keydb/keydb.go

@@ -16,67 +16,29 @@ package keydb
 
 import (
 	"context"
-	"database/sql"
+	"errors"
+	"net/url"
 
+	"github.com/matrix-org/dendrite/common/keydb/postgres"
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
-// A Database implements gomatrixserverlib.KeyDatabase and is used to store
+type Database interface {
-// the public keys for other matrix servers.
+	FetcherName() string
-type Database struct {
+	FetchKeys(ctx context.Context, requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp) (map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult, error)
-	statements serverKeyStatements
+	StoreKeys(ctx context.Context, keyMap map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult) error
 }
 
-// NewDatabase prepares a new key database.
+// NewDatabase opens a database connection.
-// It creates the necessary tables if they don't already exist.
+func NewDatabase(dataSourceName string) (Database, error) {
-// It prepares all the SQL statements that it will use.
+	uri, err := url.Parse(dataSourceName)
-// Returns an error if there was a problem talking to the database.
-func NewDatabase(dataSourceName string) (*Database, error) {
-	db, err := sql.Open("postgres", dataSourceName)
 	if err != nil {
 		return nil, err
 	}
-	d := &Database{}
+	switch uri.Scheme {
-	err = d.statements.prepare(db)
+	case "postgres":
-	if err != nil {
+		return postgres.NewDatabase(dataSourceName)
-		return nil, err
+	default:
+		return nil, errors.New("unknown schema")
 	}
-	return d, nil
-}
-
-// FetcherName implements KeyFetcher
-func (d Database) FetcherName() string {
-	return "KeyDatabase"
-}
-
-// FetchKeys implements gomatrixserverlib.KeyDatabase
-func (d *Database) FetchKeys(
-	ctx context.Context,
-	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
-) (map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult, error) {
-	return d.statements.bulkSelectServerKeys(ctx, requests)
-}
-
-// StoreKeys implements gomatrixserverlib.KeyDatabase
-func (d *Database) StoreKeys(
-	ctx context.Context,
-	keyMap map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
-) error {
-	// TODO: Inserting all the keys within a single transaction may
-	// be more efficient since the transaction overhead can be quite
-	// high for a single insert statement.
-	var lastErr error
-	for request, keys := range keyMap {
-		if err := d.statements.upsertServerKeys(ctx, request, keys); err != nil {
-			// Rather than returning immediately on error we try to insert the
-			// remaining keys.
-			// Since we are inserting the keys outside of a transaction it is
-			// possible for some of the inserts to succeed even though some
-			// of the inserts have failed.
-			// Ensuring that we always insert all the keys we can means that
-			// this behaviour won't depend on the iteration order of the map.
-			lastErr = err
-		}
-	}
-	return lastErr
 }

common/keydb/postgres/keydb.go

@@ -0,0 +1,82 @@
+// Copyright 2017 Vector Creations Ltd
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package postgres
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+// A Database implements gomatrixserverlib.KeyDatabase and is used to store
+// the public keys for other matrix servers.
+type Database struct {
+	statements serverKeyStatements
+}
+
+// NewDatabase prepares a new key database.
+// It creates the necessary tables if they don't already exist.
+// It prepares all the SQL statements that it will use.
+// Returns an error if there was a problem talking to the database.
+func NewDatabase(dataSourceName string) (*Database, error) {
+	db, err := sql.Open("postgres", dataSourceName)
+	if err != nil {
+		return nil, err
+	}
+	d := &Database{}
+	err = d.statements.prepare(db)
+	if err != nil {
+		return nil, err
+	}
+	return d, nil
+}
+
+// FetcherName implements KeyFetcher
+func (d Database) FetcherName() string {
+	return "KeyDatabase"
+}
+
+// FetchKeys implements gomatrixserverlib.KeyDatabase
+func (d *Database) FetchKeys(
+	ctx context.Context,
+	requests map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.Timestamp,
+) (map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult, error) {
+	return d.statements.bulkSelectServerKeys(ctx, requests)
+}
+
+// StoreKeys implements gomatrixserverlib.KeyDatabase
+func (d *Database) StoreKeys(
+	ctx context.Context,
+	keyMap map[gomatrixserverlib.PublicKeyLookupRequest]gomatrixserverlib.PublicKeyLookupResult,
+) error {
+	// TODO: Inserting all the keys within a single transaction may
+	// be more efficient since the transaction overhead can be quite
+	// high for a single insert statement.
+	var lastErr error
+	for request, keys := range keyMap {
+		if err := d.statements.upsertServerKeys(ctx, request, keys); err != nil {
+			// Rather than returning immediately on error we try to insert the
+			// remaining keys.
+			// Since we are inserting the keys outside of a transaction it is
+			// possible for some of the inserts to succeed even though some
+			// of the inserts have failed.
+			// Ensuring that we always insert all the keys we can means that
+			// this behaviour won't depend on the iteration order of the map.
+			lastErr = err
+		}
+	}
+	return lastErr
+}

common/keydb/postgres/server_key_table.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package keydb
+package postgres
 
 import (
 	"context"