diff --git a/roomserver/storage/postgres/invite_table.go b/roomserver/storage/postgres/invite_table.go index 176c16e48..4cddfe2e9 100644 --- a/roomserver/storage/postgres/invite_table.go +++ b/roomserver/storage/postgres/invite_table.go @@ -81,12 +81,12 @@ type inviteStatements struct { updateInviteRetiredStmt *sql.Stmt } -func createInvitesTable(db *sql.DB) error { +func CreateInvitesTable(db *sql.DB) error { _, err := db.Exec(inviteSchema) return err } -func prepareInvitesTable(db *sql.DB) (tables.Invites, error) { +func PrepareInvitesTable(db *sql.DB) (tables.Invites, error) { s := &inviteStatements{} return s, sqlutil.StatementList{ @@ -127,8 +127,8 @@ func (s *inviteStatements) UpdateInviteRetired( defer internal.CloseAndLogIfError(ctx, rows, "updateInviteRetired: rows.close() failed") var eventIDs []string + var inviteEventID string for rows.Next() { - var inviteEventID string if err = rows.Scan(&inviteEventID); err != nil { return nil, err } @@ -152,9 +152,9 @@ func (s *inviteStatements) SelectInviteActiveForUserInRoom( defer internal.CloseAndLogIfError(ctx, rows, "selectInviteActiveForUserInRoom: rows.close() failed") var result []types.EventStateKeyNID var eventIDs []string + var inviteEventID string + var senderUserNID int64 for rows.Next() { - var inviteEventID string - var senderUserNID int64 if err := rows.Scan(&inviteEventID, &senderUserNID); err != nil { return nil, nil, err } diff --git a/roomserver/storage/postgres/storage.go b/roomserver/storage/postgres/storage.go index da8d25848..27b69348b 100644 --- a/roomserver/storage/postgres/storage.go +++ b/roomserver/storage/postgres/storage.go @@ -95,7 +95,7 @@ func (d *Database) create(db *sql.DB) error { if err := createRoomAliasesTable(db); err != nil { return err } - if err := createInvitesTable(db); err != nil { + if err := CreateInvitesTable(db); err != nil { return err } if err := createMembershipTable(db); err != nil { @@ -148,7 +148,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room if err != nil { return err } - invites, err := prepareInvitesTable(db) + invites, err := PrepareInvitesTable(db) if err != nil { return err } diff --git a/roomserver/storage/sqlite3/invite_table.go b/roomserver/storage/sqlite3/invite_table.go index d54d313a9..e051d63af 100644 --- a/roomserver/storage/sqlite3/invite_table.go +++ b/roomserver/storage/sqlite3/invite_table.go @@ -69,12 +69,12 @@ type inviteStatements struct { selectInvitesAboutToRetireStmt *sql.Stmt } -func createInvitesTable(db *sql.DB) error { +func CreateInvitesTable(db *sql.DB) error { _, err := db.Exec(inviteSchema) return err } -func prepareInvitesTable(db *sql.DB) (tables.Invites, error) { +func PrepareInvitesTable(db *sql.DB) (tables.Invites, error) { s := &inviteStatements{ db: db, } @@ -119,8 +119,8 @@ func (s *inviteStatements) UpdateInviteRetired( return } defer internal.CloseAndLogIfError(ctx, rows, "UpdateInviteRetired: rows.close() failed") + var inviteEventID string for rows.Next() { - var inviteEventID string if err = rows.Scan(&inviteEventID); err != nil { return } @@ -147,9 +147,9 @@ func (s *inviteStatements) SelectInviteActiveForUserInRoom( defer internal.CloseAndLogIfError(ctx, rows, "selectInviteActiveForUserInRoom: rows.close() failed") var result []types.EventStateKeyNID var eventIDs []string + var eventID string + var senderUserNID int64 for rows.Next() { - var eventID string - var senderUserNID int64 if err := rows.Scan(&eventID, &senderUserNID); err != nil { return nil, nil, err } diff --git a/roomserver/storage/sqlite3/storage.go b/roomserver/storage/sqlite3/storage.go index e6cf1a53f..810763394 100644 --- a/roomserver/storage/sqlite3/storage.go +++ b/roomserver/storage/sqlite3/storage.go @@ -104,7 +104,7 @@ func (d *Database) create(db *sql.DB) error { if err := createRoomAliasesTable(db); err != nil { return err } - if err := createInvitesTable(db); err != nil { + if err := CreateInvitesTable(db); err != nil { return err } if err := createMembershipTable(db); err != nil { @@ -157,7 +157,7 @@ func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.Room if err != nil { return err } - invites, err := prepareInvitesTable(db) + invites, err := PrepareInvitesTable(db) if err != nil { return err } diff --git a/roomserver/storage/tables/invite_table_test.go b/roomserver/storage/tables/invite_table_test.go new file mode 100644 index 000000000..a93de23da --- /dev/null +++ b/roomserver/storage/tables/invite_table_test.go @@ -0,0 +1,90 @@ +package tables_test + +import ( + "context" + "testing" + + "github.com/matrix-org/dendrite/internal/sqlutil" + "github.com/matrix-org/dendrite/roomserver/storage/postgres" + "github.com/matrix-org/dendrite/roomserver/storage/sqlite3" + "github.com/matrix-org/dendrite/roomserver/storage/tables" + "github.com/matrix-org/dendrite/roomserver/types" + "github.com/matrix-org/dendrite/setup/config" + "github.com/matrix-org/dendrite/test" + "github.com/matrix-org/util" + "github.com/stretchr/testify/assert" +) + +func mustCreateInviteTable(t *testing.T, dbType test.DBType) (tables.Invites, func()) { + t.Helper() + connStr, close := test.PrepareDBConnectionString(t, dbType) + db, err := sqlutil.Open(&config.DatabaseOptions{ + ConnectionString: config.DataSource(connStr), + }, sqlutil.NewExclusiveWriter()) + assert.NoError(t, err) + var tab tables.Invites + switch dbType { + case test.DBTypePostgres: + err = postgres.CreateInvitesTable(db) + assert.NoError(t, err) + tab, err = postgres.PrepareInvitesTable(db) + case test.DBTypeSQLite: + err = sqlite3.CreateInvitesTable(db) + assert.NoError(t, err) + tab, err = sqlite3.PrepareInvitesTable(db) + } + assert.NoError(t, err) + + return tab, close +} + +func TestInviteTable(t *testing.T) { + ctx := context.Background() + test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) { + tab, close := mustCreateInviteTable(t, dbType) + defer close() + eventID1 := util.RandomString(16) + newInvite, err := tab.InsertInviteEvent(ctx, nil, eventID1, 1, 1, 2, []byte("")) + assert.NoError(t, err) + assert.True(t, newInvite) + + // Try adding the same invite again + newInvite, err = tab.InsertInviteEvent(ctx, nil, eventID1, 1, 1, 2, []byte("")) + assert.NoError(t, err) + assert.False(t, newInvite) + + // Add another invite for this room + eventID2 := util.RandomString(16) + newInvite, err = tab.InsertInviteEvent(ctx, nil, eventID2, 1, 1, 2, []byte("")) + assert.NoError(t, err) + assert.True(t, newInvite) + + // Add another invite for a different user + eventID := util.RandomString(16) + newInvite, err = tab.InsertInviteEvent(ctx, nil, eventID, 3, 1, 2, []byte("")) + assert.NoError(t, err) + assert.True(t, newInvite) + + stateKeyNIDs, eventIDs, err := tab.SelectInviteActiveForUserInRoom(ctx, nil, 1, 1) + assert.NoError(t, err) + assert.Equal(t, []string{eventID1, eventID2}, eventIDs) + assert.Equal(t, []types.EventStateKeyNID{2, 2}, stateKeyNIDs) + + // retire the invite + retiredEventIDs, err := tab.UpdateInviteRetired(ctx, nil, 1, 1) + assert.NoError(t, err) + assert.Equal(t, []string{eventID1, eventID2}, retiredEventIDs) + + // This should now be empty + stateKeyNIDs, eventIDs, err = tab.SelectInviteActiveForUserInRoom(ctx, nil, 1, 1) + assert.NoError(t, err) + assert.Empty(t, eventIDs) + assert.Empty(t, stateKeyNIDs) + + // Non-existent targetUserNID + stateKeyNIDs, eventIDs, err = tab.SelectInviteActiveForUserInRoom(ctx, nil, 10, 1) + assert.NoError(t, err) + assert.Empty(t, stateKeyNIDs) + assert.Empty(t, eventIDs) + }) +}