From e546a7ee4fbe07a27198c644a906f7f199ea8e68 Mon Sep 17 00:00:00 2001
From: ilija42 <57732589+ilija42@users.noreply.github.com>
Date: Wed, 29 Nov 2023 19:49:11 +0100
Subject: [PATCH] Update CleanupChainTables shell cmd to specify schema
 (#11409)

* Update CleanupChainTables shell cmd to specify schema

* Update testdb naming checks to match

* Add test that checks if cleanup chain table shell cmd works

* Fix shell cmd test db checks to match CI

* Adjust CleanupChainTables test to handle heavyweight db testdb naming

* Change CleanupChainTables to account for duplicate table names
---
 core/cmd/shell_local.go      | 27 ++++++++++++++++++++-------
 core/cmd/shell_local_test.go | 20 ++++++++++++++++++++
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/core/cmd/shell_local.go b/core/cmd/shell_local.go
index ada70736f49..dc0d44f3189 100644
--- a/core/cmd/shell_local.go
+++ b/core/cmd/shell_local.go
@@ -996,23 +996,36 @@ func (s *Shell) CleanupChainTables(c *cli.Context) error {
 	if err != nil {
 		return s.errorOut(errors.Wrap(err, "error connecting to the database"))
 	}
-
 	defer db.Close()
 
-	tablesToDeleteFromQuery := `SELECT table_name FROM information_schema.columns WHERE "column_name"=$1;`
+	// some tables with evm_chain_id (mostly job specs) are in public schema
+	tablesToDeleteFromQuery := `SELECT table_name, table_schema FROM information_schema.columns WHERE "column_name"=$1;`
 	// Delete rows from each table based on the chain_id.
 	if strings.EqualFold("EVM", c.String("type")) {
-		var tables []string
-		if err = db.Select(&tables, tablesToDeleteFromQuery, "evm_chain_id"); err != nil {
+		rows, err := db.Query(tablesToDeleteFromQuery, "evm_chain_id")
+		if err != nil {
 			return err
+		} else if rows.Err() != nil {
+			return rows.Err()
 		}
-		for _, tableName := range tables {
+
+		var tablesToDeleteFrom []string
+		for rows.Next() {
+			var name string
+			var schema string
+			if err = rows.Scan(&name, &schema); err != nil {
+				return err
+			}
+			tablesToDeleteFrom = append(tablesToDeleteFrom, schema+"."+name)
+		}
+
+		for _, tableName := range tablesToDeleteFrom {
 			query := fmt.Sprintf(`DELETE FROM %s WHERE "evm_chain_id"=$1;`, tableName)
 			_, err = db.Exec(query, c.String("id"))
 			if err != nil {
-				fmt.Printf("Error deleting rows from %s: %v\n", tableName, err)
+				fmt.Printf("Error deleting rows containing evm_chain_id from %s: %v\n", tableName, err)
 			} else {
-				fmt.Printf("Rows with chain_id %s deleted from %s.\n", c.String("id"), tableName)
+				fmt.Printf("Rows with evm_chain_id %s deleted from %s.\n", c.String("id"), tableName)
 			}
 		}
 	} else {
diff --git a/core/cmd/shell_local_test.go b/core/cmd/shell_local_test.go
index 1e030fa9e9e..fac2d7f040b 100644
--- a/core/cmd/shell_local_test.go
+++ b/core/cmd/shell_local_test.go
@@ -492,3 +492,23 @@ func TestShell_RebroadcastTransactions_AddressCheck(t *testing.T) {
 		})
 	}
 }
+
+func TestShell_CleanupChainTables(t *testing.T) {
+	// Just check if it doesn't error, command itself shouldn't be changed unless major schema changes were made.
+	// It would be really hard to write a test that accounts for schema changes, so this should be enough to alarm us that something broke.
+	config, _ := heavyweight.FullTestDBV2(t, func(c *chainlink.Config, s *chainlink.Secrets) { c.Database.Dialect = dialects.Postgres })
+	client := cmd.Shell{
+		Config: config,
+		Logger: logger.TestLogger(t),
+	}
+
+	set := flag.NewFlagSet("test", 0)
+	flagSetApplyFromAction(client.CleanupChainTables, set, "")
+	require.NoError(t, set.Set("id", testutils.FixtureChainID.String()))
+	require.NoError(t, set.Set("type", "EVM"))
+	// heavyweight creates test db named chainlink_test_uid, while usual naming is chainlink_test
+	// CleanupChainTables handles test db name with chainlink_test, but because of heavyweight test db naming we have to set danger flag
+	require.NoError(t, set.Set("danger", "true"))
+	c := cli.NewContext(nil, set, nil)
+	require.NoError(t, client.CleanupChainTables(c))
+}