From 3b88b8fd3acdce78479907b2a446c23b4be98797 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Tue, 19 May 2020 13:43:31 +0100
Subject: [PATCH] Fix postgres, fix lint

---
 roomserver/internal/input_membership.go       | 21 ++++++++++---------
 .../storage/postgres/membership_table.go      | 12 +++++++++--
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/roomserver/internal/input_membership.go b/roomserver/internal/input_membership.go
index 2383ec567..6fe866788 100644
--- a/roomserver/internal/input_membership.go
+++ b/roomserver/internal/input_membership.go
@@ -57,7 +57,6 @@ func (r *RoomserverInternalAPI) updateMemberships(
 		var ae *gomatrixserverlib.Event
 		var re *gomatrixserverlib.Event
 		targetUserNID := change.EventStateKeyNID
-		targetLocal := false
 		if change.removedEventNID != 0 {
 			ev, _ := eventMap(events).lookup(change.removedEventNID)
 			if ev != nil {
@@ -70,7 +69,7 @@ func (r *RoomserverInternalAPI) updateMemberships(
 				ae = &ev.Event
 			}
 		}
-		if updates, err = r.updateMembership(updater, targetUserNID, targetLocal, re, ae, updates); err != nil {
+		if updates, err = r.updateMembership(updater, targetUserNID, re, ae, updates); err != nil {
 			return nil, err
 		}
 	}
@@ -80,7 +79,6 @@ func (r *RoomserverInternalAPI) updateMemberships(
 func (r *RoomserverInternalAPI) updateMembership(
 	updater types.RoomRecentEventsUpdater,
 	targetUserNID types.EventStateKeyNID,
-	targetLocal bool,
 	remove, add *gomatrixserverlib.Event,
 	updates []api.OutputEvent,
 ) ([]api.OutputEvent, error) {
@@ -114,13 +112,7 @@ func (r *RoomserverInternalAPI) updateMembership(
 		return updates, nil
 	}
 
-	targetLocal = false
-	if statekey := add.StateKey(); statekey != nil {
-		_, domain, _ := gomatrixserverlib.SplitID('@', *statekey)
-		targetLocal = domain == r.Cfg.Matrix.ServerName
-	}
-
-	mu, err := updater.MembershipUpdater(targetUserNID, targetLocal)
+	mu, err := updater.MembershipUpdater(targetUserNID, r.isLocalTarget(add))
 	if err != nil {
 		return nil, err
 	}
@@ -139,6 +131,15 @@ func (r *RoomserverInternalAPI) updateMembership(
 	}
 }
 
+func (r *RoomserverInternalAPI) isLocalTarget(event *gomatrixserverlib.Event) bool {
+	targetLocal := false
+	if statekey := event.StateKey(); statekey != nil {
+		_, domain, _ := gomatrixserverlib.SplitID('@', *statekey)
+		targetLocal = domain == r.Cfg.Matrix.ServerName
+	}
+	return targetLocal
+}
+
 func updateToInviteMembership(
 	mu types.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
 	roomVersion gomatrixserverlib.RoomVersion,
diff --git a/roomserver/storage/postgres/membership_table.go b/roomserver/storage/postgres/membership_table.go
index b46cd9458..820ef4e71 100644
--- a/roomserver/storage/postgres/membership_table.go
+++ b/roomserver/storage/postgres/membership_table.go
@@ -166,7 +166,13 @@ func (s *membershipStatements) selectMembershipFromRoomAndTarget(
 func (s *membershipStatements) selectMembershipsFromRoom(
 	ctx context.Context, roomNID types.RoomNID, localOnly bool,
 ) (eventNIDs []types.EventNID, err error) {
-	rows, err := s.selectMembershipsFromRoomStmt.QueryContext(ctx, roomNID)
+	var stmt *sql.Stmt
+	if localOnly {
+		stmt = s.selectLocalMembershipsFromRoomStmt
+	} else {
+		stmt = s.selectMembershipsFromRoomStmt
+	}
+	rows, err := stmt.QueryContext(ctx, roomNID)
 	if err != nil {
 		return
 	}
@@ -187,9 +193,11 @@ func (s *membershipStatements) selectMembershipsFromRoomAndMembership(
 	roomNID types.RoomNID, membership membershipState, localOnly bool,
 ) (eventNIDs []types.EventNID, err error) {
 	var rows *sql.Rows
-	stmt := s.selectMembershipsFromRoomAndMembershipStmt
+	var stmt *sql.Stmt
 	if localOnly {
 		stmt = s.selectLocalMembershipsFromRoomAndMembershipStmt
+	} else {
+		stmt = s.selectMembershipsFromRoomAndMembershipStmt
 	}
 	rows, err = stmt.QueryContext(ctx, roomNID, membership)
 	if err != nil {