Move event state key sql to a separate file

src/github.com/matrix-org/dendrite/roomserver/storage/event_state_keys_table.go

@@ -0,0 +1,98 @@
+package storage
+
+import (
+	"database/sql"
+	"github.com/lib/pq"
+	"github.com/matrix-org/dendrite/roomserver/types"
+)
+
+const eventStateKeysSchema = `
+-- Numeric versions of the event "state_key"s. State keys tend to be reused so
+-- assigning each string a numeric ID should reduce the amount of data that
+-- needs to be stored and fetched from the database.
+-- It also means that many operations can work with int64 arrays rather than
+-- string arrays which may help reduce GC pressure.
+-- Well known state keys are pre-assigned numeric IDs:
+--   1 -> "" (the empty string)
+-- Other state keys are automatically assigned numeric IDs starting from 2**16.
+-- This leaves room to add more pre-assigned numeric IDs and clearly separates
+-- the automatically assigned IDs from the pre-assigned IDs.
+CREATE SEQUENCE IF NOT EXISTS event_state_key_nid_seq START 65536;
+CREATE TABLE IF NOT EXISTS event_state_keys (
+    -- Local numeric ID for the state key.
+    event_state_key_nid BIGINT PRIMARY KEY DEFAULT nextval('event_state_key_nid_seq'),
+    event_state_key TEXT NOT NULL CONSTRAINT event_state_key_unique UNIQUE
+);
+INSERT INTO event_state_keys (event_state_key_nid, event_state_key) VALUES
+    (1, '') ON CONFLICT DO NOTHING;
+`
+
+// Same as insertEventTypeNIDSQL
+const insertEventStateKeyNIDSQL = "" +
+	"INSERT INTO event_state_keys (event_state_key) VALUES ($1)" +
+	" ON CONFLICT ON CONSTRAINT event_state_key_unique" +
+	" DO UPDATE SET event_state_key = $1" +
+	" RETURNING (event_state_key_nid)"
+
+const selectEventStateKeyNIDSQL = "" +
+	"SELECT event_state_key_nid FROM event_state_keys WHERE event_state_key = $1"
+
+// Bulk lookup from string state key to numeric ID for that state key.
+// Takes an array of strings as the query parameter.
+const bulkSelectEventStateKeyNIDSQL = "" +
+	"SELECT event_state_key, event_state_key_nid FROM event_state_keys" +
+	" WHERE event_state_key = ANY($1)"
+
+type eventStateKeyStatements struct {
+	insertEventStateKeyNIDStmt     *sql.Stmt
+	selectEventStateKeyNIDStmt     *sql.Stmt
+	bulkSelectEventStateKeyNIDStmt *sql.Stmt
+}
+
+func (s *eventStateKeyStatements) prepare(db *sql.DB) (err error) {
+	_, err = db.Exec(eventStateKeysSchema)
+	if err != nil {
+		return
+	}
+	if s.insertEventStateKeyNIDStmt, err = db.Prepare(insertEventStateKeyNIDSQL); err != nil {
+		return
+	}
+	if s.selectEventStateKeyNIDStmt, err = db.Prepare(selectEventStateKeyNIDSQL); err != nil {
+		return
+	}
+	if s.bulkSelectEventStateKeyNIDStmt, err = db.Prepare(bulkSelectEventStateKeyNIDSQL); err != nil {
+		return
+	}
+	return
+}
+
+func (s *eventStateKeyStatements) insertEventStateKeyNID(eventStateKey string) (types.EventStateKeyNID, error) {
+	var eventStateKeyNID int64
+	err := s.insertEventStateKeyNIDStmt.QueryRow(eventStateKey).Scan(&eventStateKeyNID)
+	return types.EventStateKeyNID(eventStateKeyNID), err
+}
+
+func (s *eventStateKeyStatements) selectEventStateKeyNID(eventStateKey string) (types.EventStateKeyNID, error) {
+	var eventStateKeyNID int64
+	err := s.selectEventStateKeyNIDStmt.QueryRow(eventStateKey).Scan(&eventStateKeyNID)
+	return types.EventStateKeyNID(eventStateKeyNID), err
+}
+
+func (s *eventStateKeyStatements) bulkSelectEventStateKeyNID(eventStateKeys []string) (map[string]types.EventStateKeyNID, error) {
+	rows, err := s.bulkSelectEventStateKeyNIDStmt.Query(pq.StringArray(eventStateKeys))
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	result := make(map[string]types.EventStateKeyNID, len(eventStateKeys))
+	for rows.Next() {
+		var stateKey string
+		var stateKeyNID int64
+		if err := rows.Scan(&stateKey, &stateKeyNID); err != nil {
+			return nil, err
+		}
+		result[stateKey] = types.EventStateKeyNID(stateKeyNID)
+	}
+	return result, nil
+}

src/github.com/matrix-org/dendrite/roomserver/storage/sql.go

@@ -10,9 +10,7 @@ import (
 type statements struct {
 	partitionOffsetStatements
 	eventTypeStatements
-	insertEventStateKeyNIDStmt     *sql.Stmt
+	eventStateKeyStatements
-	selectEventStateKeyNIDStmt     *sql.Stmt
-	bulkSelectEventStateKeyNIDStmt *sql.Stmt
 	insertRoomNIDStmt              *sql.Stmt
 	selectRoomNIDStmt              *sql.Stmt
 	insertEventStmt                *sql.Stmt
@@ -39,7 +37,7 @@ func (s *statements) prepare(db *sql.DB) error {
 		return err
 	}
 
-	if err = s.prepareEventStateKeys(db); err != nil {
+	if err = s.eventStateKeyStatements.prepare(db); err != nil {
 		return err
 	}
 
@@ -58,91 +56,6 @@ func (s *statements) prepare(db *sql.DB) error {
 	return nil
 }
 
-func (s *statements) prepareEventStateKeys(db *sql.DB) (err error) {
-	_, err = db.Exec(eventStateKeysSchema)
-	if err != nil {
-		return
-	}
-	if s.insertEventStateKeyNIDStmt, err = db.Prepare(insertEventStateKeyNIDSQL); err != nil {
-		return
-	}
-	if s.selectEventStateKeyNIDStmt, err = db.Prepare(selectEventStateKeyNIDSQL); err != nil {
-		return
-	}
-	if s.bulkSelectEventStateKeyNIDStmt, err = db.Prepare(bulkSelectEventStateKeyNIDSQL); err != nil {
-		return
-	}
-	return
-}
-
-const eventStateKeysSchema = `
--- Numeric versions of the event "state_key"s. State keys tend to be reused so
--- assigning each string a numeric ID should reduce the amount of data that
--- needs to be stored and fetched from the database.
--- It also means that many operations can work with int64 arrays rather than
--- string arrays which may help reduce GC pressure.
--- Well known state keys are pre-assigned numeric IDs:
---   1 -> "" (the empty string)
--- Other state keys are automatically assigned numeric IDs starting from 2**16.
--- This leaves room to add more pre-assigned numeric IDs and clearly separates
--- the automatically assigned IDs from the pre-assigned IDs.
-CREATE SEQUENCE IF NOT EXISTS event_state_key_nid_seq START 65536;
-CREATE TABLE IF NOT EXISTS event_state_keys (
-    -- Local numeric ID for the state key.
-    event_state_key_nid BIGINT PRIMARY KEY DEFAULT nextval('event_state_key_nid_seq'),
-    event_state_key TEXT NOT NULL CONSTRAINT event_state_key_unique UNIQUE
-);
-INSERT INTO event_state_keys (event_state_key_nid, event_state_key) VALUES
-    (1, '') ON CONFLICT DO NOTHING;
-`
-
-// Same as insertEventTypeNIDSQL
-const insertEventStateKeyNIDSQL = "" +
-	"INSERT INTO event_state_keys (event_state_key) VALUES ($1)" +
-	" ON CONFLICT ON CONSTRAINT event_state_key_unique" +
-	" DO UPDATE SET event_state_key = $1" +
-	" RETURNING (event_state_key_nid)"
-
-const selectEventStateKeyNIDSQL = "" +
-	"SELECT event_state_key_nid FROM event_state_keys WHERE event_state_key = $1"
-
-// Bulk lookup from string state key to numeric ID for that state key.
-// Takes an array of strings as the query parameter.
-const bulkSelectEventStateKeyNIDSQL = "" +
-	"SELECT event_state_key, event_state_key_nid FROM event_state_keys" +
-	" WHERE event_state_key = ANY($1)"
-
-func (s *statements) insertEventStateKeyNID(eventStateKey string) (types.EventStateKeyNID, error) {
-	var eventStateKeyNID int64
-	err := s.insertEventStateKeyNIDStmt.QueryRow(eventStateKey).Scan(&eventStateKeyNID)
-	return types.EventStateKeyNID(eventStateKeyNID), err
-}
-
-func (s *statements) selectEventStateKeyNID(eventStateKey string) (types.EventStateKeyNID, error) {
-	var eventStateKeyNID int64
-	err := s.selectEventStateKeyNIDStmt.QueryRow(eventStateKey).Scan(&eventStateKeyNID)
-	return types.EventStateKeyNID(eventStateKeyNID), err
-}
-
-func (s *statements) bulkSelectEventStateKeyNID(eventStateKeys []string) (map[string]types.EventStateKeyNID, error) {
-	rows, err := s.bulkSelectEventStateKeyNIDStmt.Query(pq.StringArray(eventStateKeys))
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close()
-
-	result := make(map[string]types.EventStateKeyNID, len(eventStateKeys))
-	for rows.Next() {
-		var stateKey string
-		var stateKeyNID int64
-		if err := rows.Scan(&stateKey, &stateKeyNID); err != nil {
-			return nil, err
-		}
-		result[stateKey] = types.EventStateKeyNID(stateKeyNID)
-	}
-	return result, nil
-}
-
 func (s *statements) prepareRooms(db *sql.DB) (err error) {
 	_, err = db.Exec(roomsSchema)
 	if err != nil {