From 21e61636e136cef6e4e2074ff5c34ffd32cfda7d Mon Sep 17 00:00:00 2001
From: Kegan Dougal <kegan@matrix.org>
Date: Tue, 3 Nov 2020 15:31:10 +0000
Subject: [PATCH] Add nodes table for event metadata

---
 internal/mscs/msc2836/storage.go | 42 ++++++++++++++++++++++++++++----
 1 file changed, 37 insertions(+), 5 deletions(-)

diff --git a/internal/mscs/msc2836/storage.go b/internal/mscs/msc2836/storage.go
index 4fd7a614b..4e975054a 100644
--- a/internal/mscs/msc2836/storage.go
+++ b/internal/mscs/msc2836/storage.go
@@ -12,14 +12,18 @@ import (
 
 type Database interface {
 	// StoreRelation stores the parent->child and child->parent relationship for later querying.
+	// Also stores the event metadata e.g timestamp
 	StoreRelation(ctx context.Context, ev *gomatrixserverlib.HeaderedEvent) error
+	// ChildrenForParent returns the event IDs of events who have the given `eventID` as an m.relationship with the
+	// provided `relType`.
 	ChildrenForParent(ctx context.Context, eventID, relType string) ([]string, error)
 }
 
 type DB struct {
 	db                          *sql.DB
 	writer                      sqlutil.Writer
-	insertRelationStmt          *sql.Stmt
+	insertEdgeStmt              *sql.Stmt
+	insertNodeStmt              *sql.Stmt
 	selectChildrenForParentStmt *sql.Stmt
 }
 
@@ -46,15 +50,26 @@ func newPostgresDatabase(dbOpts *config.DatabaseOptions) (Database, error) {
 		rel_type TEXT NOT NULL,
 		CONSTRAINT msc2836_edges UNIQUE (parent_event_id, child_event_id, rel_type)
 	);
+
+	CREATE TABLE IF NOT EXISTS msc2836_nodes (
+		event_id TEXT PRIMARY KEY NOT NULL,
+		origin_server_ts BIGINT NOT NULL,
+		room_id TEXT NOT NULL
+	);
 	`)
 	if err != nil {
 		return nil, err
 	}
-	if d.insertRelationStmt, err = d.db.Prepare(`
+	if d.insertEdgeStmt, err = d.db.Prepare(`
 		INSERT INTO msc2836_edges(parent_event_id, child_event_id, rel_type) VALUES($1, $2, $3) ON CONFLICT DO NOTHING
 	`); err != nil {
 		return nil, err
 	}
+	if d.insertNodeStmt, err = d.db.Prepare(`
+		INSERT INTO msc2836_nodes(event_id, origin_server_ts, room_id) VALUES($1, $2, $3) ON CONFLICT DO NOTHING
+	`); err != nil {
+		return nil, err
+	}
 	if d.selectChildrenForParentStmt, err = d.db.Prepare(`
 		SELECT child_event_id FROM msc2836_edges WHERE parent_event_id = $1 AND rel_type = $2
 	`); err != nil {
@@ -78,15 +93,26 @@ func newSQLiteDatabase(dbOpts *config.DatabaseOptions) (Database, error) {
 		rel_type TEXT NOT NULL,
 		UNIQUE (parent_event_id, child_event_id, rel_type)
 	);
+
+	CREATE TABLE IF NOT EXISTS msc2836_nodes (
+		event_id TEXT PRIMARY KEY NOT NULL,
+		origin_server_ts BIGINT NOT NULL,
+		room_id TEXT NOT NULL
+	);
 	`)
 	if err != nil {
 		return nil, err
 	}
-	if d.insertRelationStmt, err = d.db.Prepare(`
+	if d.insertEdgeStmt, err = d.db.Prepare(`
 		INSERT INTO msc2836_edges(parent_event_id, child_event_id, rel_type) VALUES($1, $2, $3) ON CONFLICT (parent_event_id, child_event_id, rel_type) DO NOTHING
 	`); err != nil {
 		return nil, err
 	}
+	if d.insertNodeStmt, err = d.db.Prepare(`
+		INSERT INTO msc2836_nodes(event_id, origin_server_ts, room_id) VALUES($1, $2, $3) ON CONFLICT DO NOTHING
+	`); err != nil {
+		return nil, err
+	}
 	if d.selectChildrenForParentStmt, err = d.db.Prepare(`
 		SELECT child_event_id FROM msc2836_edges WHERE parent_event_id = $1 AND rel_type = $2
 	`); err != nil {
@@ -100,8 +126,14 @@ func (p *DB) StoreRelation(ctx context.Context, ev *gomatrixserverlib.HeaderedEv
 	if parent == "" || child == "" {
 		return nil
 	}
-	_, err := p.insertRelationStmt.ExecContext(ctx, parent, child, relType)
-	return err
+	return p.writer.Do(p.db, nil, func(txn *sql.Tx) error {
+		_, err := txn.Stmt(p.insertEdgeStmt).ExecContext(ctx, parent, child, relType)
+		if err != nil {
+			return err
+		}
+		_, err = txn.Stmt(p.insertNodeStmt).ExecContext(ctx, ev.EventID(), ev.OriginServerTS(), ev.RoomID())
+		return err
+	})
 }
 
 func (p *DB) ChildrenForParent(ctx context.Context, eventID, relType string) ([]string, error) {