From 35c14addd541a98435a474c0d13b04bf1300e7fd Mon Sep 17 00:00:00 2001
From: wwestgarth <william@vega.xyz>
Date: Wed, 25 Sep 2024 09:57:40 +0100
Subject: [PATCH] feat: add check on data node startup that DB schema version
 is what it should be

---
 cmd/data-node/commands/start/node_pre.go |  6 ++++
 datanode/sqlstore/sqlstore.go            | 38 ++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/cmd/data-node/commands/start/node_pre.go b/cmd/data-node/commands/start/node_pre.go
index 5e1528e350e..9d035ac6f58 100644
--- a/cmd/data-node/commands/start/node_pre.go
+++ b/cmd/data-node/commands/start/node_pre.go
@@ -181,6 +181,12 @@ func (l *NodeCommand) persistentPre([]string) (err error) {
 		return fmt.Errorf("failed to apply data retention policies:%w", err)
 	}
 
+	// check that the schema version matches the latest migration, because if it doesn't queries might fail if rows/tables
+	// it expects to exist don't
+	if err := sqlstore.CheckSchemaVersionsSynced(l.Log, conf.SQLStore.ConnectionConfig, sqlstore.EmbedMigrations); err != nil {
+		return err
+	}
+
 	preLog.Info("Enabling SQL stores")
 
 	l.transactionalConnectionSource, err = sqlstore.NewTransactionalConnectionSource(l.ctx, preLog, l.conf.SQLStore.ConnectionConfig)
diff --git a/datanode/sqlstore/sqlstore.go b/datanode/sqlstore/sqlstore.go
index 7a8e8dde280..f22b72075b6 100644
--- a/datanode/sqlstore/sqlstore.go
+++ b/datanode/sqlstore/sqlstore.go
@@ -205,6 +205,44 @@ func RevertToSchemaVersionZero(log *logging.Logger, config ConnectionConfig, fs
 	return nil
 }
 
+// CheckSchemaVersionsSynced checks if if the current migrated version of the DB matches the latest available version.
+// If they are not in sync then we are running code expecting a particular version, but the rows/columns we ask for
+// might not exist.
+func CheckSchemaVersionsSynced(log *logging.Logger, config ConnectionConfig, fs fs.FS) error {
+	goose.SetBaseFS(fs)
+	goose.SetLogger(log.GooseLogger())
+	goose.SetVerbose(true)
+
+	log.Info("Checking database version matches code version")
+	poolConfig, err := config.GetPoolConfig()
+	if err != nil {
+		return fmt.Errorf("failed to get pool config: %w", err)
+	}
+
+	db := stdlib.OpenDB(*poolConfig.ConnConfig)
+	defer db.Close()
+
+	current, err := goose.GetDBVersion(db)
+	if err != nil {
+		return err
+	}
+
+	migrations, err := goose.CollectMigrations(SQLMigrationsDir, current, current+1)
+	if err != nil {
+		return err
+	}
+
+	if migrations.Len() != 0 {
+		last, err := migrations.Last()
+		if err != nil {
+			return err
+		}
+		return fmt.Errorf("Schema is version %d, latest should be %d", current, last.Version)
+	}
+
+	return nil
+}
+
 func WipeDatabaseAndMigrateSchemaToVersion(log *logging.Logger, config ConnectionConfig, version int64, fs fs.FS, verbose bool) error {
 	log = log.Named("db-wipe-migrate")
 	goose.SetBaseFS(fs)