Fix sqlite tests and close rows

This commit is contained in:
Kegan Dougal 2020-07-15 18:05:49 +01:00
parent ad840c54b9
commit 83105d4c34
2 changed files with 6 additions and 9 deletions

View file

@ -19,6 +19,7 @@ import (
"database/sql" "database/sql"
"time" "time"
"github.com/matrix-org/dendrite/internal"
"github.com/matrix-org/dendrite/internal/sqlutil" "github.com/matrix-org/dendrite/internal/sqlutil"
"github.com/matrix-org/dendrite/keyserver/api" "github.com/matrix-org/dendrite/keyserver/api"
"github.com/matrix-org/dendrite/keyserver/storage/tables" "github.com/matrix-org/dendrite/keyserver/storage/tables"
@ -108,6 +109,7 @@ func (s *deviceKeysStatements) SelectBatchDeviceKeys(ctx context.Context, userID
if err != nil { if err != nil {
return nil, err return nil, err
} }
defer internal.CloseAndLogIfError(ctx, rows, "selectBatchDeviceKeysStmt: rows.close() failed")
deviceIDMap := make(map[string]bool) deviceIDMap := make(map[string]bool)
for _, d := range deviceIDs { for _, d := range deviceIDs {
deviceIDMap[d] = true deviceIDMap[d] = true

View file

@ -17,9 +17,9 @@ package sqlite3
import ( import (
"context" "context"
"database/sql" "database/sql"
"strings"
"time" "time"
"github.com/matrix-org/dendrite/internal"
"github.com/matrix-org/dendrite/internal/sqlutil" "github.com/matrix-org/dendrite/internal/sqlutil"
"github.com/matrix-org/dendrite/keyserver/api" "github.com/matrix-org/dendrite/keyserver/api"
"github.com/matrix-org/dendrite/keyserver/storage/tables" "github.com/matrix-org/dendrite/keyserver/storage/tables"
@ -47,7 +47,7 @@ const selectDeviceKeysSQL = "" +
"SELECT key_json FROM keyserver_device_keys WHERE user_id=$1 AND device_id=$2" "SELECT key_json FROM keyserver_device_keys WHERE user_id=$1 AND device_id=$2"
const selectBatchDeviceKeysSQL = "" + const selectBatchDeviceKeysSQL = "" +
"SELECT device_id, key_json FROM keyserver_device_keys WHERE user_id=$1 AND device_id IN ($2)" "SELECT device_id, key_json FROM keyserver_device_keys WHERE user_id=$1"
type deviceKeysStatements struct { type deviceKeysStatements struct {
db *sql.DB db *sql.DB
@ -81,16 +81,11 @@ func (s *deviceKeysStatements) SelectBatchDeviceKeys(ctx context.Context, userID
for _, d := range deviceIDs { for _, d := range deviceIDs {
deviceIDMap[d] = true deviceIDMap[d] = true
} }
iDeviceIDs := make([]interface{}, len(deviceIDs)+1) rows, err := s.selectBatchDeviceKeysStmt.QueryContext(ctx, userID)
iDeviceIDs[0] = userID
for i := range deviceIDs {
iDeviceIDs[i+1] = deviceIDs[i]
}
querySQL := strings.Replace(selectBatchDeviceKeysSQL, "($2)", sqlutil.QueryVariadic(len(deviceIDs)), 1)
rows, err := s.db.QueryContext(ctx, querySQL, iDeviceIDs...)
if err != nil { if err != nil {
return nil, err return nil, err
} }
defer internal.CloseAndLogIfError(ctx, rows, "selectBatchDeviceKeysStmt: rows.close() failed")
var result []api.DeviceKeys var result []api.DeviceKeys
for rows.Next() { for rows.Next() {
var dk api.DeviceKeys var dk api.DeviceKeys