make the migration commitable to master

cmd/util/cmd/execution-state-extract/cmd.go

@@ -325,7 +325,7 @@ func run(*cobra.Command, []string) {
 	}
 
 	// Validate chain ID
-	_ = flow.ChainID(flagChain).Chain()
+	chain := flow.ChainID(flagChain).Chain()
 
 	if flagNoReport {
 		log.Warn().Msgf("--no-report flag is deprecated")
@@ -422,6 +422,8 @@ func run(*cobra.Command, []string) {
 		var migs []migrations.NamedMigration
 
 		switch flagMigration {
+		case "add-keys":
+			migs = append(migs, addKeysMigration(log.Logger, flagOutputDir, flagNWorker, chain.ChainID())...)
 		default:
 			log.Fatal().Msgf("unknown migration: %s", flagMigration)
 		}

cmd/util/cmd/execution-state-extract/execution_state_extract.go

@@ -10,7 +10,9 @@ import (
 	"github.com/rs/zerolog/log"
 	"golang.org/x/sync/errgroup"
 
+	"github.com/onflow/flow-go-sdk/crypto"
 	migrators "github.com/onflow/flow-go/cmd/util/ledger/migrations"
+	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/ledger/common/hash"
@@ -356,3 +358,37 @@ func createTrieFromPayloads(logger zerolog.Logger, payloads []*ledger.Payload) (
 
 	return newTrie, nil
 }
+
+func addKeysMigration(
+	log zerolog.Logger,
+	outputDir string,
+	workerCount int,
+	chainID flow.ChainID,
+) []migrators.NamedMigration {
+
+	log.Info().Msg("initializing add-keys migrations ...")
+
+	rwf := reporters.NewReportFileWriterFactory(outputDir, log)
+
+	key, err := crypto.DecodePublicKeyHex(crypto.ECDSA_P256, "711d4cd9930d695ef5c79b668d321f92ba00ed8280fded52c0fa2b15501411d026fe6fb4be3ec894facd3a00f04e32e2db5f5696d3b2b3419e4fba89fb95dca8")
+	if err != nil {
+		panic("failed to decode key")
+	}
+
+	namedMigrations := []migrators.NamedMigration{
+		{
+			Name: "account-usage-migration",
+			Migrate: migrators.NewAccountBasedMigration(
+				log,
+				workerCount,
+				[]migrators.AccountBasedMigration{
+					migrators.NewAddKeyMigration(chainID, key, rwf),
+				},
+			),
+		},
+	}
+
+	log.Info().Msg("initialized migrations")
+
+	return namedMigrations
+}

cmd/util/ledger/migrations/add_key_migration.go

@@ -4,20 +4,25 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
+	"github.com/rs/zerolog"
 
-	"github.com/onflow/cadence/runtime/common"
+	"github.com/onflow/cadence/common"
 	"github.com/onflow/crypto/hash"
-	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go-sdk/crypto"
-
+	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
 	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
 	"github.com/onflow/flow-go/fvm"
 	"github.com/onflow/flow-go/fvm/systemcontracts"
 	"github.com/onflow/flow-go/model/flow"
 )
 
+// This migration is not safe to run on actual networks.
+// It is used for getting control over system accounts so that they can be tested
+// in a copied environment. When we need to do this it is best to create a branch
+// where this variable is set to true, and use that temporary branc to run migrations.
+const IAmSureIWantToRunThisMigration = false
+
 // AddKeyMigration adds a new key to the core contracts accounts
 type AddKeyMigration struct {
 	log zerolog.Logger
@@ -35,6 +40,9 @@ func NewAddKeyMigration(
 	key crypto.PublicKey,
 	rwf reporters.ReportWriterFactory,
 ) *AddKeyMigration {
+	if !IAmSureIWantToRunThisMigration {
+		panic("Cannot run AddKeyMigration migration")
+	}
 
 	addresses := make(map[common.Address]AddKeyMigrationAccountPublicKeyData)
 	sc := systemcontracts.SystemContractsForChain(chainID).All()
@@ -91,6 +99,9 @@ func (m *AddKeyMigration) MigrateAccount(
 	address common.Address,
 	accountRegisters *registers.AccountRegisters,
 ) error {
+	if !IAmSureIWantToRunThisMigration {
+		panic("Cannot run AddKeyMigration migration")
+	}
 
 	keyData, ok := m.accountsToAddKeyTo[address]
 	if !ok {

cmd/util/ledger/migrations/add_key_migration_test.go

@@ -1,92 +1,13 @@
 package migrations
 
 import (
-	"context"
-	"crypto/rand"
 	"testing"
-
-	"github.com/onflow/cadence/runtime/common"
-	"github.com/rs/zerolog"
-	"github.com/stretchr/testify/require"
-
-	"github.com/onflow/flow-go-sdk/crypto"
-
-	"github.com/onflow/flow-go/cmd/util/ledger/util"
-	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
-	"github.com/onflow/flow-go/fvm"
-	"github.com/onflow/flow-go/fvm/systemcontracts"
-	"github.com/onflow/flow-go/model/flow"
 )
 
-func TestCoreContractsKeys(t *testing.T) {
+func Test_fail_if_migration_enabled(t *testing.T) {
 	t.Parallel()
-
-	log := zerolog.New(zerolog.NewTestWriter(t))
-
-	// Get the old payloads
-	payloads, err := util.PayloadsFromEmulatorSnapshot(snapshotPath)
-	require.NoError(t, err)
-
-	registersByAccount, err := registers.NewByAccountFromPayloads(payloads)
-	require.NoError(t, err)
-
-	chainID := flow.Emulator
-	sc := systemcontracts.SystemContractsForChain(chainID)
-	serviceAccountAddress := sc.FlowServiceAccount.Address
-
-	serviceRegisters := registersByAccount.AccountRegisters(string(serviceAccountAddress.Bytes()))
-
-	pk, err := crypto.GeneratePrivateKey(crypto.ECDSA_P256, makeSeed(t))
-	require.NoError(t, err)
-	expectedKey := pk.PublicKey()
-	rwf := &testReportWriterFactory{}
-
-	mig := NewAddKeyMigration(
-		chainID,
-		expectedKey,
-		rwf,
-	)
-	defer func() {
-		err := mig.Close()
-		require.NoError(t, err)
-	}()
-
-	err = mig.InitMigration(log, registersByAccount, 1)
-	require.NoError(t, err)
-
-	ctx := context.Background()
-	err = mig.MigrateAccount(ctx, common.Address(serviceAccountAddress), serviceRegisters)
-	require.NoError(t, err)
-
-	// Create all the runtime components we need for the migration
-	migrationRuntime, err := NewInterpreterMigrationRuntime(
-		serviceRegisters,
-		chainID,
-		InterpreterMigrationRuntimeConfig{},
-	)
-	require.NoError(t, err)
-
-	// The last key should be the one we added
-	keys, err := migrationRuntime.Accounts.GetPublicKeyCount(serviceAccountAddress)
-	require.NoError(t, err)
-
-	key, err := migrationRuntime.Accounts.GetPublicKey(serviceAccountAddress, keys-1)
-	require.NoError(t, err)
-
-	require.Equal(t, expectedKey.String(), key.PublicKey.String())
-	require.Equal(t, fvm.AccountKeyWeightThreshold, key.Weight)
-}
-
-func Test_DO_NOT_MERGE(t *testing.T) {
-	t.Parallel()
-	// this branch should not be merged to master
-	// This is only to be used for migration mainnet testing
-	t.Fail()
-}
-
-func makeSeed(t *testing.T) []byte {
-	seed := make([]byte, 32)
-	_, err := rand.Read(seed)
-	require.NoError(t, err)
-	return seed
+	// prevent merging this to master branch if enabled
+	if IAmSureIWantToRunThisMigration {
+		t.Fail()
+	}
 }