From 1d900e66cd8f677cdbfa6237b630862dabf5004f Mon Sep 17 00:00:00 2001
From: Brandur <brandur@brandur.org>
Date: Mon, 6 Apr 2026 15:25:14 -0500
Subject: [PATCH] When getting tables to truncate, panic on unknown version

This one's related to code review feedback here [1]. When getting a list
of tables that need to be truncated for the `TestSchema` helpers, it was
possible to receive migration versions that didn't exist and the
function would tolerate it, returning the latest known list of tables.
This often works, but even where it does there's some risk of a new
table coming in that doesn't end up getting cleaned up properly.

Here, panic when receiving an unknown version number so that the problem
can be corrected immediately. The panic may seem undesirable, but these
functions are only used in tests, so it should be fine.

[1] https://github.com/riverqueue/riverpro/pull/281#discussion_r3040295326
---
 riverdriver/river_driver_interface.go      |  7 +++++--
 riverdriver/river_driver_interface_test.go | 13 +++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/riverdriver/river_driver_interface.go b/riverdriver/river_driver_interface.go
index aed92b7e..43eacb26 100644
--- a/riverdriver/river_driver_interface.go
+++ b/riverdriver/river_driver_interface.go
@@ -15,6 +15,7 @@ package riverdriver
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io/fs"
 	"time"
 
@@ -865,6 +866,7 @@ type TableTruncateParams struct {
 //
 // API is not stable. DO NOT USE.
 func MigrationLineMainTruncateTables(version int) []string {
+	// 0 value must be handled and should always point to latest migration version
 	switch version {
 	case 1:
 		return nil // don't truncate `river_migrate`
@@ -872,8 +874,9 @@ func MigrationLineMainTruncateTables(version int) []string {
 		return []string{"river_job", "river_leader"}
 	case 4:
 		return []string{"river_job", "river_leader", "river_queue"}
+	case 0, 5, 6:
+		return []string{"river_job", "river_leader", "river_queue", "river_client", "river_client_queue"}
 	}
 
-	// 0 (zero value), 5, 6
-	return []string{"river_job", "river_leader", "river_queue", "river_client", "river_client_queue"}
+	panic(fmt.Sprintf("unrecognized migration version: %d", version))
 }
diff --git a/riverdriver/river_driver_interface_test.go b/riverdriver/river_driver_interface_test.go
index 866d122b..978e299c 100644
--- a/riverdriver/river_driver_interface_test.go
+++ b/riverdriver/river_driver_interface_test.go
@@ -250,6 +250,19 @@ func TestJobSetStateSnoozed(t *testing.T) { //nolint:dupl
 	})
 }
 
+func TestMigrationLineMainTruncateTables(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ZeroValueDoesNotPanic", func(t *testing.T) {
+		t.Parallel()
+
+		require.NotPanics(t, func() {
+			tables := MigrationLineMainTruncateTables(0)
+			require.NotEmpty(t, tables)
+		})
+	})
+}
+
 func TestJobSetStateSnoozedAvailable(t *testing.T) { //nolint:dupl
 	t.Parallel()