[release-12.3.2] Fix primary keys (#115884)

Convert unique keys in 3 tables to primary keys

pkg/services/sqlstore/migrations/cloud_migrations.go

@@ -198,10 +198,11 @@ func addCloudMigrationsMigrations(mg *Migrator) {
 		Postgres("ALTER TABLE cloud_migration_resource ALTER COLUMN resource_uid TYPE VARCHAR(255);"))
 
 	mg.AddMigration("create cloud_migration_snapshot_partition table v1", NewAddTableMigration(migrationSnapshotPartitionTable))
-	mg.AddMigration("add cloud_migration_snapshot_partition srp_unique index", NewAddIndexMigration(migrationSnapshotPartitionTable, &Index{
+	srpUniqueIndex := Index{
 		Name: "srp_unique",
 		Cols: []string{"snapshot_uid", "resource_type", "partition_number"}, Type: UniqueIndex,
-	}))
+	}
+	mg.AddMigration("add cloud_migration_snapshot_partition srp_unique index", NewAddIndexMigration(migrationSnapshotPartitionTable, &srpUniqueIndex))
 	mg.AddMigration("add resource_storage_type column to cloud_migration_snapshot table", NewAddColumnMigration(migrationSnapshotTable, &Column{
 		Name:     "resource_storage_type",
 		Type:     DB_Varchar,
@@ -224,4 +225,16 @@ func addCloudMigrationsMigrations(mg *Migrator) {
 		Type:     DB_Blob,
 		Nullable: true,
 	}))
+
+	updatedCloudMigrationSnapshotPartitionTable := Table{
+		Name: "cloud_migration_snapshot_partition",
+		Columns: []*Column{
+			{Name: "snapshot_uid", Type: DB_NVarchar, Length: 40, Nullable: false, IsPrimaryKey: true},
+			{Name: "partition_number", Type: DB_Int, Nullable: false, IsPrimaryKey: true},
+			{Name: "resource_type", Type: DB_Varchar, Length: 255, Nullable: false, IsPrimaryKey: true},
+			{Name: "data", Type: DB_LongBlob, Nullable: false},
+		},
+		PrimaryKeys: []string{"snapshot_uid", "resource_type", "partition_number"},
+	}
+	ConvertUniqueKeyToPrimaryKey(mg, srpUniqueIndex, updatedCloudMigrationSnapshotPartitionTable)
 }

pkg/services/sqlstore/migrator/dialect.go

@@ -253,7 +253,7 @@ func (b *BaseDialect) CopyTableData(sourceTable string, targetTable string, sour
 	targetColsSQL := b.QuoteColList(targetCols)
 
 	quote := b.dialect.Quote
-	return fmt.Sprintf("INSERT INTO %s (%s) SELECT %s FROM %s", quote(targetTable), targetColsSQL, sourceColsSQL, quote(sourceTable))
+	return fmt.Sprintf("INSERT INTO %s (%s)\nSELECT %s\nFROM %s", quote(targetTable), targetColsSQL, sourceColsSQL, quote(sourceTable))
 }
 
 func (b *BaseDialect) DropTable(tableName string) string {

pkg/services/sqlstore/migrator/migrations.go

@@ -1,6 +1,8 @@
 package migrator
 
 import (
+	"fmt"
+	"slices"
 	"strings"
 )
 
@@ -271,3 +273,155 @@ func NewTableCharsetMigration(tableName string, columns []*Column) *TableCharset
 func (m *TableCharsetMigration) SQL(d Dialect) string {
 	return d.UpdateTableSQL(m.tableName, m.columns)
 }
+
+type addPrimaryKeyMigration struct {
+	MigrationBase
+	tableName string
+	uniqueKey Index
+
+	// Used for Sqlite recreation of the table. Temporary table will have tableName + "_new" suffix.
+	table Table
+}
+
+func (m *addPrimaryKeyMigration) SQL(d Dialect) string {
+	if d.DriverName() == SQLite {
+		// Final SQL will do following in the individual statements:
+		// 1. Create new temporary table
+		// 2. Copy data from old table to temporary table
+		// 3. Drop old table, rename temporary table to original name
+		// 4. Recreate indexes for table.
+		//
+		// For example:
+		//
+		//	CREATE TABLE file_new
+		//	(
+		//		path                    TEXT     NOT NULL,
+		//		path_hash               TEXT     NOT NULL,
+		//		parent_folder_path_hash TEXT     NOT NULL,
+		//		contents                BLOB     NOT NULL,
+		//		etag                    TEXT     NOT NULL,
+		//		cache_control           TEXT     NOT NULL,
+		//		content_disposition     TEXT     NOT NULL,
+		//		updated                 DATETIME NOT NULL,
+		//		created                 DATETIME NOT NULL,
+		//		size                    INTEGER  NOT NULL,
+		//		mime_type               TEXT     NOT NULL,
+		//
+		//		PRIMARY KEY (path_hash)
+		//	);
+		//
+		//	INSERT INTO file_new (path, path_hash, parent_folder_path_hash, contents, etag, cache_control, content_disposition, updated, created, size, mime_type)
+		//	SELECT path, path_hash, parent_folder_path_hash, contents, etag, cache_control, content_disposition, updated, created, size, mime_type FROM file;
+		//
+		//	DROP TABLE file;
+		//	ALTER TABLE file_new RENAME TO file;
+		//
+		//	CREATE INDEX IDX_file_parent_folder_path_hash ON file (parent_folder_path_hash);
+
+		tempTable := m.table
+		tempTable.Name = m.tableName + "_new"
+
+		statements := strings.Builder{}
+
+		statements.WriteString(d.CreateTableSQL(&tempTable))
+		statements.WriteString("\n") // CreateTableSQL adds semicolon
+
+		cols := make([]string, 0, len(tempTable.Columns))
+		for _, col := range tempTable.Columns {
+			cols = append(cols, col.Name)
+		}
+		statements.WriteString(d.CopyTableData(m.tableName, tempTable.Name, cols, cols))
+		statements.WriteString(";\n")
+
+		statements.WriteString(d.DropTable(m.tableName))
+		statements.WriteString(";\n")
+
+		statements.WriteString(d.RenameTable(tempTable.Name, m.tableName))
+		statements.WriteString(";\n")
+
+		for _, idx := range tempTable.Indices {
+			// Use real table name, not temporary one now
+			statements.WriteString(d.CreateIndexSQL(m.tableName, idx))
+			statements.WriteString("\n") // CreateIndexSQL adds semicolon
+		}
+
+		return statements.String()
+	} else if d.DriverName() == Postgres {
+		quotesCols := make([]string, 0, len(m.uniqueKey.Cols))
+		for _, c := range m.uniqueKey.Cols {
+			quotesCols = append(quotesCols, d.Quote(c))
+		}
+
+		return fmt.Sprintf(`
+		DO $$
+		BEGIN
+			-- Drop the unique constraint if it exists
+			DROP INDEX IF EXISTS %s;
+
+			-- Add primary key if it doesn't already exist
+			IF NOT EXISTS (SELECT 1 FROM pg_index i WHERE indrelid = '%s'::regclass AND indisprimary) THEN
+				ALTER TABLE %s ADD PRIMARY KEY (%s);
+			END IF;
+		END $$;`, d.Quote(m.uniqueKey.XName(m.tableName)), m.tableName, d.Quote(m.tableName), strings.Join(quotesCols, ","))
+	} else {
+		return ""
+	}
+}
+
+// ConvertUniqueKeyToPrimaryKey adds series of migrations to convert existing unique key to PRIMARY KEY.
+// For Sqlite this means recreating the table, which only works if there are no foreign keys referencing the table.
+func ConvertUniqueKeyToPrimaryKey(mg *Migrator, uniqueKey Index, finalTable Table) {
+	tableName := finalTable.Name
+	if tableName == "" {
+		panic("invalid table name")
+	}
+	if len(uniqueKey.Cols) == 0 || uniqueKey.Type != UniqueIndex {
+		panic("invalid unique type")
+	}
+	if !slices.Equal(uniqueKey.Cols, finalTable.PrimaryKeys) {
+		panic("invalid primary key in the final table")
+	}
+
+	colPks := map[string]bool{}
+	for _, col := range finalTable.Columns {
+		if col.IsPrimaryKey {
+			colPks[col.Name] = true
+		}
+	}
+	for _, c := range uniqueKey.Cols {
+		if !colPks[c] {
+			panic(fmt.Sprintf("column %s is not part of primary key in the table definition", c))
+		}
+	}
+
+	columnsList := strings.Join(uniqueKey.Cols, ",")
+
+	mysqlQuote := NewDialect(MySQL).Quote
+	mysqlQuotedColumns := make([]string, 0, len(uniqueKey.Cols))
+	for _, col := range uniqueKey.Cols {
+		mysqlQuotedColumns = append(mysqlQuotedColumns, mysqlQuote(col))
+	}
+
+	// migration 1 is to handle cases where the table was created with sql_generate_invisible_primary_key = ON
+	// in this case we need to do the conversion in one sql statement
+	mysqlMigration1 := NewRawSQLMigration("").Mysql(fmt.Sprintf(`
+	  ALTER TABLE %s
+	  DROP PRIMARY KEY,
+	  DROP COLUMN my_row_id,
+	  DROP INDEX %s,
+	  ADD PRIMARY KEY (%s);
+	`, tableName, uniqueKey.XName(tableName), strings.Join(mysqlQuotedColumns, ",")))
+	mysqlMigration1.Condition = &IfColumnExistsCondition{TableName: tableName, ColumnName: "my_row_id"}
+	mg.AddMigration(fmt.Sprintf("drop my_row_id and add primary key with columns %s to table %s if my_row_id exists (auto-generated mysql column)", columnsList, tableName), mysqlMigration1)
+
+	mysqlMigration2 := NewRawSQLMigration("").Mysql(fmt.Sprintf(`ALTER TABLE %s DROP INDEX %s`, tableName, uniqueKey.XName(tableName)))
+	mysqlMigration2.Condition = &IfIndexExistsCondition{TableName: tableName, IndexName: uniqueKey.XName(tableName)}
+	mg.AddMigration(fmt.Sprintf("drop unique index %s from %s table if it exists (mysql)", uniqueKey.XName(tableName), tableName), mysqlMigration2)
+
+	mysqlMigration3 := NewRawSQLMigration("").Mysql(fmt.Sprintf(`ALTER TABLE %s ADD PRIMARY KEY (%s)`, tableName, strings.Join(mysqlQuotedColumns, ",")))
+	mysqlMigration3.Condition = &IfPrimaryKeyNotExistsCondition{TableName: tableName}
+	mg.AddMigration(fmt.Sprintf("add primary key with columns %s to table %s if it doesn't exist (mysql)", columnsList, tableName), mysqlMigration3)
+
+	// postgres and sqlite statements are idempotent so we can have only one condition-less migration
+	mg.AddMigration(fmt.Sprintf("add primary key with columns %s to table %s (postgres and sqlite)", columnsList, tableName), &addPrimaryKeyMigration{tableName: tableName, uniqueKey: uniqueKey, table: finalTable})
+}

pkg/services/sqlstore/migrator/migrations_test.go

@@ -0,0 +1,79 @@
+package migrator
+
+import (
+	_ "embed"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+//go:embed testdata/sqlite_file_migration_statement.sql
+var sqliteMigrationStatement string
+
+func TestConvertUniqueKeyToPrimaryKey(t *testing.T) {
+	names := []string{
+		"drop my_row_id and add primary key with columns path_hash,etag to table file if my_row_id exists (auto-generated mysql column)",
+		"drop unique index UQE_file_path_hash_etag from file table if it exists (mysql)",
+		"add primary key with columns path_hash,etag to table file if it doesn't exist (mysql)",
+		"add primary key with columns path_hash,etag to table file (postgres and sqlite)",
+	}
+	expectedMigrations := map[string][]ExpectedMigration{
+		MySQL: {
+			{Id: names[0], SQL: `
+			  ALTER TABLE file
+			  DROP PRIMARY KEY,
+			  DROP COLUMN my_row_id,
+			  DROP INDEX UQE_file_path_hash_etag,
+			  ADD PRIMARY KEY (` + "`path_hash`" + `,` + "`etag`" + `);`},
+			{Id: names[1], SQL: "ALTER TABLE file DROP INDEX UQE_file_path_hash_etag"},
+			{Id: names[2], SQL: "ALTER TABLE file ADD PRIMARY KEY (`path_hash`,`etag`)"},
+			{Id: names[3], SQL: ""},
+		},
+		Postgres: {
+			{Id: names[0], SQL: ""},
+			{Id: names[1], SQL: ""},
+			{Id: names[2], SQL: ""},
+			{Id: names[3], SQL: `
+				DO $$
+				BEGIN
+					-- Drop the unique constraint if it exists
+					DROP INDEX IF EXISTS "UQE_file_path_hash_etag";
+
+					-- Add primary key if it doesn't already exist
+					IF NOT EXISTS (SELECT 1 FROM pg_index i WHERE indrelid = 'file'::regclass AND indisprimary) THEN
+					ALTER TABLE "file" ADD PRIMARY KEY ("path_hash","etag");
+				END IF;
+				END $$;`},
+		},
+		SQLite: {
+			{Id: names[0], SQL: ""},
+			{Id: names[1], SQL: ""},
+			{Id: names[2], SQL: ""},
+			{Id: names[3], SQL: sqliteMigrationStatement}, // Embed used here because sqlite statement is full of backquotes.
+		},
+	}
+
+	for dialectName, migrations := range expectedMigrations {
+		t.Run(dialectName, func(t *testing.T) {
+			err := CheckExpectedMigrations(dialectName, migrations, func(migrator *Migrator) {
+				ConvertUniqueKeyToPrimaryKey(migrator,
+					Index{Cols: []string{"path_hash", "etag"}, Type: UniqueIndex}, // Convert this unique key to primary key
+					Table{
+						Name: "file",
+						Columns: []*Column{
+							{Name: "path", Type: DB_NVarchar, Length: 1024, Nullable: false},
+							{Name: "path_hash", Type: DB_NVarchar, Length: 64, Nullable: false, IsPrimaryKey: true},
+							{Name: "parent_folder_path_hash", Type: DB_NVarchar, Length: 64, Nullable: false},
+							{Name: "contents", Type: DB_Blob, Nullable: false},
+							{Name: "etag", Type: DB_NVarchar, Length: 32, Nullable: false, IsPrimaryKey: true},
+						},
+						PrimaryKeys: []string{"path_hash", "etag"},
+						Indices: []*Index{
+							{Cols: []string{"parent_folder_path_hash"}},
+						},
+					})
+			})
+			require.NoError(t, err)
+		})
+	}
+}

pkg/services/sqlstore/migrator/migrator.go

@@ -8,7 +8,6 @@ import (
 
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/golang-migrate/migrate/v4/database"
-	"github.com/grafana/grafana/pkg/util/sqlite"
 	_ "github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel"
@@ -17,6 +16,8 @@ import (
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/atomic"
 
+	"github.com/grafana/grafana/pkg/util/sqlite"
+
 	"github.com/grafana/grafana/pkg/util/xorm"
 
 	"github.com/grafana/grafana/pkg/infra/log"
@@ -67,12 +68,16 @@ func NewMigrator(engine *xorm.Engine, cfg *setting.Cfg) *Migrator {
 
 // NewScopedMigrator should only be used for the transition to a new storage engine
 func NewScopedMigrator(engine *xorm.Engine, cfg *setting.Cfg, scope string) *Migrator {
+	return newMigrator(engine, cfg, scope, NewDialect(engine.DriverName()))
+}
+
+func newMigrator(engine *xorm.Engine, cfg *setting.Cfg, scope string, dialect Dialect) *Migrator {
 	mg := &Migrator{
 		Cfg:          cfg,
 		DBEngine:     engine,
 		migrations:   make([]Migration, 0),
 		migrationIds: make(map[string]struct{}),
-		Dialect:      NewDialect(engine.DriverName()),
+		Dialect:      dialect,
 		metrics: migratorMetrics{
 			migCount: prometheus.NewCounterVec(prometheus.CounterOpts{
 				Namespace: "grafana_database",

pkg/services/sqlstore/migrator/testdata/sqlite_file_migration_statement.sql

@@ -0,0 +1,23 @@
+CREATE TABLE IF NOT EXISTS `file_new` (
+    `path` TEXT NOT NULL
+    , `path_hash` TEXT NOT NULL
+    , `parent_folder_path_hash` TEXT NOT NULL
+    , `contents` BLOB NOT NULL
+    , `etag` TEXT NOT NULL
+    , PRIMARY KEY ( `path_hash`,`etag` ));
+
+INSERT INTO `file_new` (`path`
+  , `path_hash`
+  , `parent_folder_path_hash`
+  , `contents`
+  , `etag`)
+SELECT `path`
+  , `path_hash`
+  , `parent_folder_path_hash`
+  , `contents`
+  , `etag`
+FROM `file`;
+
+DROP TABLE IF EXISTS `file`;
+ALTER TABLE `file_new` RENAME TO `file`;
+CREATE INDEX `IDX_file_parent_folder_path_hash` ON `file` (`parent_folder_path_hash`);

pkg/services/sqlstore/migrator/testing.go

@@ -0,0 +1,51 @@
+package migrator
+
+import (
+	"fmt"
+	"strings"
+)
+
+type ExpectedMigration struct {
+	Id  string
+	SQL string
+}
+
+// CheckExpectedMigrations verifies that given migrations exist in migrator after running addMigrations function,
+// that they are in the same order and have expected SQL.
+func CheckExpectedMigrations(dialectName string, expected []ExpectedMigration, addMigrations func(migrator *Migrator)) error {
+	d := NewDialect(dialectName)
+	mg := newMigrator(nil, nil, "", d)
+	addMigrations(mg)
+
+	migrations := mg.migrations
+	migrationNames := make([]string, 0, len(migrations))
+	for _, m := range expected {
+		for ; len(migrations) > 0 && migrations[0].Id() != m.Id; migrations = migrations[1:] {
+			migrationNames = append(migrationNames, migrations[0].Id())
+		}
+
+		if len(migrations) == 0 {
+			return fmt.Errorf("migration `%s` not found, existing migrations:\n%s", m.Id, strings.Join(migrationNames, "\n"))
+		}
+
+		sql := migrations[0].SQL(d)
+		if normalizeLines(m.SQL) != normalizeLines(sql) {
+			return fmt.Errorf("migration `%s` has wrong SQL:\nexpected:\n%s\nactual:\n%s", m.Id, m.SQL, sql)
+		}
+	}
+	return nil
+}
+
+func normalizeLines(sql string) string {
+	lines := strings.Split(sql, "\n")
+	result := strings.Builder{}
+	for _, l := range lines {
+		l := strings.TrimSpace(l)
+		if l == "" {
+			continue
+		}
+		result.WriteString(l)
+		result.WriteString("\n")
+	}
+	return result.String()
+}

pkg/storage/secret/migrator/migrator.go

@@ -211,4 +211,23 @@ func (*SecretDB) AddMigration(mg *migrator.Migrator) {
 	mg.AddMigration("add data_key_id index to "+TableNameEncryptedValue, migrator.NewAddIndexMigration(encryptedValueTable, &migrator.Index{
 		Cols: []string{"data_key_id"},
 	}))
+
+	encryptedValueTableUniqueKey := migrator.Index{Cols: []string{"namespace", "name", "version"}, Type: migrator.UniqueIndex}
+	updatedEncryptedValueTable := migrator.Table{
+		Name: TableNameEncryptedValue,
+		Columns: []*migrator.Column{
+			{Name: "namespace", Type: migrator.DB_NVarchar, Length: 253, Nullable: false, IsPrimaryKey: true}, // Limit enforced by K8s.
+			{Name: "name", Type: migrator.DB_NVarchar, Length: 253, Nullable: false, IsPrimaryKey: true},
+			{Name: "version", Type: migrator.DB_BigInt, Nullable: false, IsPrimaryKey: true},
+			{Name: "encrypted_data", Type: migrator.DB_Blob, Nullable: false},
+			{Name: "created", Type: migrator.DB_BigInt, Nullable: false},
+			{Name: "updated", Type: migrator.DB_BigInt, Nullable: false},
+			{Name: "data_key_id", Type: migrator.DB_NVarchar, Length: 100, Nullable: false, Default: "''"},
+		},
+		PrimaryKeys: []string{"namespace", "name", "version"},
+		Indices: []*migrator.Index{
+			{Cols: []string{"data_key_id"}},
+		},
+	}
+	migrator.ConvertUniqueKeyToPrimaryKey(mg, encryptedValueTableUniqueKey, updatedEncryptedValueTable)
 }

pkg/storage/unified/sql/db/migrations/resource_mig.go

@@ -185,5 +185,18 @@ func initResourceTables(mg *migrator.Migrator) string {
 		Name: "UQE_resource_last_import_time_last_import_time",
 	}))
 
+	oldResourceVersionUniqueKey := migrator.Index{Cols: []string{"group", "resource"}, Type: migrator.UniqueIndex}
+	updatedResourceVersionTable := migrator.Table{
+		Name: "resource_version",
+		Columns: []*migrator.Column{
+			{Name: "group", Type: migrator.DB_NVarchar, Length: 190, Nullable: false, IsPrimaryKey: true},
+			{Name: "resource", Type: migrator.DB_NVarchar, Length: 190, Nullable: false, IsPrimaryKey: true},
+			{Name: "resource_version", Type: migrator.DB_BigInt, Nullable: false},
+		},
+		PrimaryKeys: []string{"group", "resource"},
+	}
+
+	migrator.ConvertUniqueKeyToPrimaryKey(mg, oldResourceVersionUniqueKey, updatedResourceVersionTable)
+
 	return marker
 }