From 77732e106d058d070da5958ef51544e3539b7307 Mon Sep 17 00:00:00 2001
From: Andy Yang <yang@cockroachlabs.com>
Date: Fri, 13 Jan 2023 14:54:59 -0500
Subject: [PATCH] upgrades: modify InjectLegacyTable to handle dynamically
 assigned IDs

This patch modifies the `InjectLegacyTable` function, which is used by
upgrade tests to inject an old version of a table descriptor, to handle
the case where the table has a dynamically assigned ID.

Release note: None
---
 pkg/upgrade/upgrades/BUILD.bazel     |  1 +
 pkg/upgrade/upgrades/helpers_test.go | 27 ++++++++++++++++++++++-----
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/pkg/upgrade/upgrades/BUILD.bazel b/pkg/upgrade/upgrades/BUILD.bazel
index 0a2cbb247ade..6f3231c3ab55 100644
--- a/pkg/upgrade/upgrades/BUILD.bazel
+++ b/pkg/upgrade/upgrades/BUILD.bazel
@@ -130,6 +130,7 @@ go_test(
         "//pkg/sql/catalog/descbuilder",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/catalog/descs",
+        "//pkg/sql/catalog/schemadesc",
         "//pkg/sql/catalog/schematelemetry/schematelemetrycontroller",
         "//pkg/sql/catalog/systemschema",
         "//pkg/sql/catalog/tabledesc",
diff --git a/pkg/upgrade/upgrades/helpers_test.go b/pkg/upgrade/upgrades/helpers_test.go
index 1a3e86d67346..e5be578d47f9 100644
--- a/pkg/upgrade/upgrades/helpers_test.go
+++ b/pkg/upgrade/upgrades/helpers_test.go
@@ -20,6 +20,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemadesc"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/stretchr/testify/assert"
@@ -76,12 +78,27 @@ func InjectLegacyTable(
 	err := s.InternalDB().(descs.DB).DescsTxn(ctx, func(
 		ctx context.Context, txn descs.Txn,
 	) error {
-		id := table.GetID()
-		tab, err := txn.Descriptors().MutableByID(txn.KV()).Table(ctx, id)
-		if err != nil {
-			return err
+		deprecatedDesc := getDeprecatedDescriptor()
+		var tab *tabledesc.Mutable
+		switch id := table.GetID(); id {
+		// If the table descriptor does not have a valid ID, it must be a system
+		// table with a dynamically-allocated ID.
+		case descpb.InvalidID:
+			var err error
+			tab, err = txn.Descriptors().MutableByName(txn.KV()).Table(ctx,
+				systemschema.SystemDB, schemadesc.GetPublicSchema(), table.GetName())
+			if err != nil {
+				return err
+			}
+			deprecatedDesc.ID = tab.GetID()
+		default:
+			var err error
+			tab, err = txn.Descriptors().MutableByID(txn.KV()).Table(ctx, id)
+			if err != nil {
+				return err
+			}
 		}
-		builder := tabledesc.NewBuilder(getDeprecatedDescriptor())
+		builder := tabledesc.NewBuilder(deprecatedDesc)
 		if err := builder.RunPostDeserializationChanges(); err != nil {
 			return err
 		}