| // The fsmigrator executable migrates various data from firestore to an SQL database. |
| // It uses port forwarding, as that is the simplest approach and there shouldn't be |
| // too much data. |
| package main |
| |
| import ( |
| "context" |
| "crypto/md5" |
| "flag" |
| "strings" |
| "sync/atomic" |
| "time" |
| |
| "cloud.google.com/go/firestore" |
| "github.com/cockroachdb/cockroach-go/v2/crdb/crdbpgx" |
| "github.com/google/uuid" |
| "github.com/jackc/pgx/v4" |
| "github.com/jackc/pgx/v4/pgxpool" |
| "golang.org/x/sync/errgroup" |
| "google.golang.org/api/iterator" |
| |
| ifirestore "go.skia.org/infra/go/firestore" |
| "go.skia.org/infra/go/skerr" |
| "go.skia.org/infra/go/sklog" |
| "go.skia.org/infra/go/util" |
| "go.skia.org/infra/golden/go/expectations" |
| "go.skia.org/infra/golden/go/fs_utils" |
| "go.skia.org/infra/golden/go/sql" |
| "go.skia.org/infra/golden/go/sql/schema" |
| "go.skia.org/infra/golden/go/types" |
| ) |
| |
| func main() { |
| var ( |
| fsProjectID = flag.String("fs_project_id", "skia-firestore", "The project with the firestore instance. Datastore and Firestore can't be in the same project.") |
| oldFSNamespace = flag.String("old_fs_namespace", "", "Typically the instance id. e.g. 'chrome-gpu', 'skia', etc") |
| newSQLDatabase = flag.String("new_sql_db", "", "Something like the instance id (no dashes)") |
| ) |
| flag.Parse() |
| |
| if *oldFSNamespace == "" { |
| sklog.Fatalf("You must include fs_namespace") |
| } |
| |
| if *newSQLDatabase == "" { |
| sklog.Fatalf("You must include new_sql_db") |
| } |
| |
| ctx := context.Background() |
| fsClient, err := ifirestore.NewClient(ctx, *fsProjectID, "gold", *oldFSNamespace, nil) |
| if err != nil { |
| sklog.Fatalf("Unable to configure Firestore: %s", err) |
| } |
| |
| u := sql.GetConnectionURL("root@localhost:26234", *newSQLDatabase) |
| conf, err := pgxpool.ParseConfig(u) |
| if err != nil { |
| sklog.Fatalf("error getting postgres config %s: %s", u, err) |
| } |
| conf.MaxConns = 16 |
| db, err := pgxpool.ConnectConfig(ctx, conf) |
| if err != nil { |
| sklog.Info("You must run\nkubectl port-forward gold-cockroachdb-0 26234:26234") |
| sklog.Fatalf("error connecting to the database: %s", err) |
| } |
| |
| oldNamespace := v3Impl{client: fsClient} |
| |
| // We need to look up the groupings by test name, because the old groupings were just the |
| // test name and not combined with the corpus. As such, we fetch all the groupings from the SQL, |
| // which has been filled through ingesting. |
| nameToGroupings, err := fetchGroupings(ctx, db) |
| if err != nil { |
| sklog.Fatalf("Getting groupings from SQL: %s", err) |
| } |
| sklog.Infof("Fetched %d groupings from SQL DB", len(nameToGroupings)) |
| |
| records, err := oldNamespace.fetchTriageRecords(ctx) |
| if err != nil { |
| sklog.Fatalf("Fetching triage records: %s", err) |
| } |
| sklog.Infof("Should migrate %d records", len(records)) |
| |
| if err := storeAndCombineTriageRecords(ctx, db, records); err != nil { |
| sklog.Fatalf("storing triage records %s", err) |
| } |
| |
| sklog.Infof("Fetching deltas") |
| deltas, err := oldNamespace.fetchExpectationDeltas(ctx) |
| if err != nil { |
| sklog.Fatalf("Fetching expectation deltas: %s", err) |
| } |
| |
| if err := storeExpectationDeltas(ctx, db, deltas, nameToGroupings); err != nil { |
| sklog.Fatalf("storing expectation deltas: %s", err) |
| } |
| |
| sklog.Infof("Fetching expectations") |
| exp, err := oldNamespace.fetchExpectations(ctx) |
| if err != nil { |
| sklog.Fatalf("fetching expectations %s", err) |
| } |
| |
| // pass in deltas so we can link in the triage record to the expectations. |
| unowned, err := storeExpectations(ctx, db, exp, nameToGroupings, deltas) |
| if err != nil { |
| sklog.Fatalf("storing expectation deltas: %s", err) |
| } |
| |
| // Write the catchall expectation record. If somehow an expectation exists, but wasn't covered |
| // by a delta (e.g. broken old data), we assign it to this catchall record. |
| _, err = db.Exec(ctx, `UPSERT INTO ExpectationRecords |
| (expectation_record_id, user_name, triage_time, num_changes) VALUES ($1, $2, $3, $4)`, |
| catchAllUUID, "sql_migrator", time.Now(), unowned) |
| if err != nil { |
| sklog.Fatalf("creating catch-all record: %s", err) |
| } |
| |
| sklog.Info("done") |
| } |
| |
| func fetchGroupings(ctx context.Context, db *pgxpool.Pool) (map[types.TestName][]schema.GroupingID, error) { |
| rv := map[types.TestName][]schema.GroupingID{} |
| rows, err := db.Query(ctx, `SELECT keys -> 'name', grouping_id FROM Groupings`) |
| if err != nil { |
| return nil, skerr.Wrap(err) |
| } |
| defer rows.Close() |
| for rows.Next() { |
| var name types.TestName |
| var gID schema.GroupingID |
| if err := rows.Scan(&name, &gID); err != nil { |
| return nil, skerr.Wrap(err) |
| } |
| rv[name] = append(rv[name], gID) |
| } |
| count := 0 |
| for _, groupings := range rv { |
| if len(groupings) > 1 { |
| count++ |
| } |
| } |
| sklog.Infof("%d out of %d test names belonged to multiple corpora", count, len(rv)) |
| return rv, nil |
| } |
| |
| func storeAndCombineTriageRecords(ctx context.Context, db *pgxpool.Pool, toStore map[string][]v3TriageRecord) error { |
| sklog.Infof("have %d partitions", len(toStore)) |
| const batchSize = 1000 |
| eg, ctx := errgroup.WithContext(ctx) |
| for b, r := range toStore { |
| branchName, records := b, r |
| sklog.Infof("Writing records from partition %s", branchName) |
| eg.Go(func() error { |
| return util.ChunkIter(len(records), batchSize, func(startIdx int, endIdx int) error { |
| if err := ctx.Err(); err != nil { |
| return skerr.Wrap(err) |
| } |
| batch := records[startIdx:endIdx] |
| statement := `INSERT INTO ExpectationRecords |
| (expectation_record_id, user_name, triage_time, num_changes, branch_name) VALUES ` |
| const valuesPerRow = 5 |
| arguments := make([]interface{}, 0, valuesPerRow*len(batch)) |
| for _, record := range batch { |
| // We can turn the old IDs into a UUID by hashing the bytes. This is faster than |
| // having to return the new random UUIDs. |
| newID := uuid.Must(uuid.FromBytes(hash(record.ID))) |
| arguments = append(arguments, newID, record.UserName, record.TS, record.Changes) |
| if branchName == v3PrimaryPartition { |
| arguments = append(arguments, nil) |
| } else { |
| arguments = append(arguments, branchName) |
| } |
| } |
| statement += sql.ValuesPlaceholders(valuesPerRow, len(batch)) |
| statement += `ON CONFLICT DO NOTHING` |
| err := crdbpgx.ExecuteTx(ctx, db, pgx.TxOptions{}, func(tx pgx.Tx) error { |
| _, err := tx.Exec(ctx, statement, arguments...) |
| return err // don't wrap - might get retried |
| }) |
| return skerr.Wrap(err) |
| }) |
| }) |
| } |
| if err := eg.Wait(); err != nil { |
| return skerr.Wrapf(err, "writing records to SQL") |
| } |
| return nil |
| } |
| |
| func storeExpectationDeltas(ctx context.Context, db *pgxpool.Pool, toStore map[string][]v3ExpectationChange, nameToGroupings map[types.TestName][]schema.GroupingID) error { |
| sklog.Infof("have %d partitions", len(toStore)) |
| const batchSize = 1000 |
| eg, ctx := errgroup.WithContext(ctx) |
| for b, d := range toStore { |
| branchName, deltas := b, d |
| sklog.Infof("Writing deltas from partition %s", branchName) |
| eg.Go(func() error { |
| return util.ChunkIter(len(deltas), batchSize, func(startIdx int, endIdx int) error { |
| if err := ctx.Err(); err != nil { |
| return skerr.Wrap(err) |
| } |
| batch := deltas[startIdx:endIdx] |
| statement := `INSERT INTO ExpectationDeltas |
| (expectation_record_id, grouping_id, digest, label_before, label_after) VALUES ` |
| const valuesPerRow = 5 |
| arguments := make([]interface{}, 0, valuesPerRow*len(batch)) |
| for _, delta := range batch { |
| newID := uuid.Must(uuid.FromBytes(hash(delta.RecordID))) |
| groupings, ok := nameToGroupings[delta.Grouping] |
| if !ok { |
| sklog.Warningf("Unknown grouping for name %s on branch %s", delta.Grouping, branchName) |
| continue |
| } |
| dBytes, err := sql.DigestToBytes(delta.Digest) |
| if err != nil { |
| sklog.Warningf("Corrupt digest %q on branch %s", delta.Digest, branchName) |
| continue |
| } |
| // Some groupings were on multiple corpora, so we need to store the record for |
| // all of them. |
| for _, gID := range groupings { |
| arguments = append(arguments, |
| newID, gID, dBytes, convertLabel(delta.LabelBefore), convertLabel(delta.AffectedRange.Label)) |
| } |
| } |
| if len(arguments) == 0 { |
| return nil |
| } |
| // Need to divide here to account for skipped rows. |
| statement += sql.ValuesPlaceholders(valuesPerRow, len(arguments)/valuesPerRow) |
| statement += `ON CONFLICT DO NOTHING` |
| err := crdbpgx.ExecuteTx(ctx, db, pgx.TxOptions{}, func(tx pgx.Tx) error { |
| _, err := tx.Exec(ctx, statement, arguments...) |
| return err // don't wrap - might get retried |
| }) |
| return skerr.Wrap(err) |
| }) |
| }) |
| } |
| if err := eg.Wait(); err != nil { |
| return skerr.Wrapf(err, "writing expectations to SQL") |
| } |
| return nil |
| } |
| |
| var ( |
| catchAllUUID = uuid.MustParse("00000000-0000-0000-0000-000000000000") |
| ) |
| |
| func storeExpectations(ctx context.Context, db *pgxpool.Pool, toStore map[string][]v3ExpectationEntry, nameToGroupings map[types.TestName][]schema.GroupingID, deltas map[string][]v3ExpectationChange) (int, error) { |
| sklog.Infof("have %d partitions", len(toStore)) |
| extraExpectations := int32(0) |
| eg, ctx := errgroup.WithContext(ctx) |
| eg.Go(func() error { |
| extra, err := storePrimaryBranchExpectations(ctx, db, toStore[v3PrimaryPartition], nameToGroupings, deltas[v3PrimaryPartition]) |
| atomic.AddInt32(&extraExpectations, int32(extra)) |
| return skerr.Wrap(err) |
| }) |
| |
| for b, e := range toStore { |
| branchName, exps := b, e |
| if branchName == v3PrimaryPartition { |
| continue |
| } |
| sklog.Infof("Writing expectations from partition %s", branchName) |
| eg.Go(func() error { |
| extra, err := storeSecondaryBranchExpectations(ctx, db, branchName, exps, nameToGroupings, deltas[branchName]) |
| atomic.AddInt32(&extraExpectations, int32(extra)) |
| return skerr.Wrap(err) |
| }) |
| } |
| if err := eg.Wait(); err != nil { |
| return 0, skerr.Wrapf(err, "writing deltas to SQL") |
| } |
| return int(extraExpectations), nil |
| } |
| |
| func storePrimaryBranchExpectations(ctx context.Context, db *pgxpool.Pool, exps []v3ExpectationEntry, nameToGroupings map[types.TestName][]schema.GroupingID, deltas []v3ExpectationChange) (int, error) { |
| const batchSize = 1000 |
| extra := 0 |
| return extra, util.ChunkIter(len(exps), batchSize, func(startIdx int, endIdx int) error { |
| if err := ctx.Err(); err != nil { |
| return skerr.Wrap(err) |
| } |
| batch := exps[startIdx:endIdx] |
| statement := `UPSERT INTO Expectations |
| (grouping_id, digest, label, expectation_record_id) VALUES ` |
| const valuesPerRow = 4 |
| arguments := make([]interface{}, 0, valuesPerRow*len(batch)) |
| for _, exp := range batch { |
| groupings, ok := nameToGroupings[exp.Grouping] |
| if !ok { |
| continue |
| } |
| dBytes, err := sql.DigestToBytes(exp.Digest) |
| if err != nil { |
| sklog.Warningf("Corrupt digest %q on branch %s", exp.Digest, v3PrimaryPartition) |
| continue |
| } |
| label := exp.Ranges[0].Label |
| recordID, ok := find(exp.Grouping, exp.Digest, label, deltas) |
| // some test names correspond to more than one grouping. We need to write the |
| // expectations for both. |
| for _, gID := range groupings { |
| if !ok { |
| extra++ |
| } |
| arguments = append(arguments, gID, dBytes, convertLabel(label), recordID) |
| } |
| } |
| if len(arguments) == 0 { |
| return nil |
| } |
| // Need to divide here to account for skipped rows. |
| statement += sql.ValuesPlaceholders(valuesPerRow, len(arguments)/valuesPerRow) |
| err := crdbpgx.ExecuteTx(ctx, db, pgx.TxOptions{}, func(tx pgx.Tx) error { |
| _, err := tx.Exec(ctx, statement, arguments...) |
| return err // don't wrap - might get retried |
| }) |
| return skerr.Wrap(err) |
| }) |
| } |
| |
| func storeSecondaryBranchExpectations(ctx context.Context, db *pgxpool.Pool, branchName string, exps []v3ExpectationEntry, nameToGroupings map[types.TestName][]schema.GroupingID, deltas []v3ExpectationChange) (int, error) { |
| const batchSize = 1000 |
| extra := 0 |
| return extra, util.ChunkIter(len(exps), batchSize, func(startIdx int, endIdx int) error { |
| if err := ctx.Err(); err != nil { |
| return skerr.Wrap(err) |
| } |
| batch := exps[startIdx:endIdx] |
| statement := `UPSERT INTO SecondaryBranchExpectations |
| (branch_name, grouping_id, digest, label, expectation_record_id) VALUES ` |
| const valuesPerRow = 5 |
| arguments := make([]interface{}, 0, valuesPerRow*len(batch)) |
| for _, exp := range batch { |
| groupings, ok := nameToGroupings[exp.Grouping] |
| if !ok { |
| continue |
| } |
| dBytes, err := sql.DigestToBytes(exp.Digest) |
| if err != nil { |
| sklog.Warningf("Corrupt digest %q on branch %s", exp.Digest, branchName) |
| continue |
| } |
| label := exp.Ranges[0].Label |
| recordID, ok := find(exp.Grouping, exp.Digest, label, deltas) |
| // account for test names that belong to multiple corpora. |
| for _, gID := range groupings { |
| if !ok { |
| extra++ |
| } |
| arguments = append(arguments, branchName, gID, dBytes, convertLabel(label), recordID) |
| } |
| } |
| if len(arguments) == 0 { |
| return nil |
| } |
| // Need to divide here to account for skipped rows. |
| statement += sql.ValuesPlaceholders(valuesPerRow, len(arguments)/valuesPerRow) |
| err := crdbpgx.ExecuteTx(ctx, db, pgx.TxOptions{}, func(tx pgx.Tx) error { |
| _, err := tx.Exec(ctx, statement, arguments...) |
| return err // don't wrap - might get retried |
| }) |
| return skerr.Wrap(err) |
| }) |
| } |
| |
| // find returns a matching record id for the given change and true or the catch-all UUID and false. |
| func find(grouping types.TestName, digest types.Digest, label expectations.LabelInt, deltas []v3ExpectationChange) (uuid.UUID, bool) { |
| for _, delta := range deltas { |
| if delta.Grouping == grouping && delta.Digest == digest && delta.AffectedRange.Label == label { |
| return uuid.Must(uuid.FromBytes(hash(delta.RecordID))), true |
| } |
| } |
| return catchAllUUID, false |
| } |
| |
| func convertLabel(label expectations.LabelInt) schema.ExpectationLabel { |
| switch label { |
| case expectations.UntriagedInt: |
| return schema.LabelUntriaged |
| case expectations.PositiveInt: |
| return schema.LabelPositive |
| case expectations.NegativeInt: |
| return schema.LabelNegative |
| } |
| return schema.LabelUntriaged |
| } |
| |
| func hash(id string) []byte { |
| h := md5.Sum([]byte(id)) |
| return h[:] |
| } |
| |
| const ( |
| v3Partitions = "expstore_partitions_v3" |
| v3ExpectationEntries = "entries" |
| v3RecordEntries = "triage_records" |
| v3ChangeEntries = "triage_changes" |
| |
| v3PrimaryPartition = "master" |
| |
| v3DigestField = "digest" |
| |
| maxRetries = 10 |
| maxOperationTime = 10 * time.Minute |
| ) |
| |
| type v3Impl struct { |
| client *ifirestore.Client |
| } |
| |
| type v3ExpectationEntry struct { |
| Grouping types.TestName `firestore:"grouping"` |
| Digest types.Digest `firestore:"digest"` |
| Updated time.Time `firestore:"updated"` |
| LastUsed time.Time `firestore:"last_used"` |
| Ranges []v3TriageRange `firestore:"ranges"` |
| NeedsGC bool `firestore:"needs_gc"` |
| } |
| |
| func (e *v3ExpectationEntry) id() string { |
| s := string(e.Grouping) + "|" + string(e.Digest) |
| // firestore gets cranky if there are / in key names |
| return strings.Replace(s, "/", "-", -1) |
| } |
| |
| type v3TriageRange struct { |
| FirstIndex int `firestore:"first_index"` |
| LastIndex int `firestore:"last_index"` |
| Label expectations.LabelInt `firestore:"label"` |
| } |
| |
| type v3TriageRecord struct { |
| ID string |
| UserName string `firestore:"user"` |
| TS time.Time `firestore:"ts"` |
| Changes int `firestore:"changes"` |
| Committed bool `firestore:"committed"` |
| } |
| |
| type v3ExpectationChange struct { |
| // RecordID refers to a document in the records collection. |
| RecordID string `firestore:"record_id"` |
| Grouping types.TestName `firestore:"grouping"` |
| Digest types.Digest `firestore:"digest"` |
| AffectedRange v3TriageRange `firestore:"affected_range"` |
| LabelBefore expectations.LabelInt `firestore:"label_before"` |
| } |
| |
| func (v v3Impl) fetchTriageRecords(ctx context.Context) (map[string][]v3TriageRecord, error) { |
| // maps partition to records |
| rv := map[string][]v3TriageRecord{} |
| partitionIterator := v.client.Collection(v3Partitions).DocumentRefs(ctx) |
| p, err := partitionIterator.Next() |
| for ; err == nil; p, err = partitionIterator.Next() { |
| var records []v3TriageRecord |
| partition := p.ID |
| sklog.Infof("Partition %s", partition) |
| recordIterator := v.client.Collection(v3Partitions).Doc(partition).Collection(v3RecordEntries).Documents(ctx) |
| docs, err := recordIterator.GetAll() |
| if err != nil { |
| return nil, skerr.Wrapf(err, "getting records for %s", partition) |
| } |
| for _, doc := range docs { |
| var r v3TriageRecord |
| if err := doc.DataTo(&r); err != nil { |
| sklog.Warning("Corrupt triage record with id %s", doc.Ref.ID) |
| continue |
| } |
| r.ID = doc.Ref.ID |
| records = append(records, r) |
| } |
| rv[partition] = records |
| } |
| if err != iterator.Done { |
| return nil, skerr.Wrap(err) |
| } |
| return rv, nil |
| } |
| |
| func (v v3Impl) fetchExpectationDeltas(ctx context.Context) (map[string][]v3ExpectationChange, error) { |
| rv := map[string][]v3ExpectationChange{} // Maps partition -> entries |
| |
| partitionIterator := v.client.Collection(v3Partitions).DocumentRefs(ctx) |
| p, err := partitionIterator.Next() |
| for ; err == nil; p, err = partitionIterator.Next() { |
| partition := p.ID |
| const numShards = 16 |
| base := v.client.Collection(v3Partitions).Doc(partition).Collection(v3ChangeEntries) |
| |
| queries := fs_utils.ShardOnDigest(base, v3DigestField, numShards) |
| shardedEntries := make([][]v3ExpectationChange, numShards) |
| |
| err := v.client.IterDocsInParallel(ctx, "loadExpectationDeltas", partition, queries, maxRetries, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error { |
| if doc == nil { |
| return nil |
| } |
| entry := v3ExpectationChange{} |
| if err := doc.DataTo(&entry); err != nil { |
| id := doc.Ref.ID |
| return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal expectationChange with id %s", id) |
| } |
| |
| shardedEntries[i] = append(shardedEntries[i], entry) |
| return nil |
| }) |
| |
| if err != nil { |
| return nil, skerr.Wrapf(err, "fetching expectation deltas for partition %s", partition) |
| } |
| |
| var combinedEntries []v3ExpectationChange |
| for _, shard := range shardedEntries { |
| combinedEntries = append(combinedEntries, shard...) |
| } |
| |
| rv[partition] = combinedEntries |
| } |
| return rv, nil |
| } |
| |
| func (v v3Impl) fetchExpectations(ctx context.Context) (map[string][]v3ExpectationEntry, error) { |
| rv := map[string][]v3ExpectationEntry{} // Maps partition -> entries |
| |
| partitionIterator := v.client.Collection(v3Partitions).DocumentRefs(ctx) |
| p, err := partitionIterator.Next() |
| for ; err == nil; p, err = partitionIterator.Next() { |
| partition := p.ID |
| |
| const numShards = 16 |
| base := v.client.Collection(v3Partitions).Doc(partition).Collection(v3ExpectationEntries) |
| queries := fs_utils.ShardOnDigest(base, v3DigestField, numShards) |
| shardedEntries := make([][]v3ExpectationEntry, numShards) |
| |
| err := v.client.IterDocsInParallel(ctx, "loadExpectations", partition, queries, maxRetries, maxOperationTime, func(i int, doc *firestore.DocumentSnapshot) error { |
| if doc == nil { |
| return nil |
| } |
| entry := v3ExpectationEntry{} |
| if err := doc.DataTo(&entry); err != nil { |
| id := doc.Ref.ID |
| return skerr.Wrapf(err, "corrupt data in firestore, could not unmarshal expectationEntry with id %s", id) |
| } |
| if len(entry.Ranges) == 0 { |
| // This should never happen, but we'll ignore these malformed entries if they do. |
| return nil |
| } |
| shardedEntries[i] = append(shardedEntries[i], entry) |
| return nil |
| }) |
| |
| if err != nil { |
| return nil, skerr.Wrapf(err, "fetching expectations for partition %s", partition) |
| } |
| |
| var combinedEntries []v3ExpectationEntry |
| for _, shard := range shardedEntries { |
| combinedEntries = append(combinedEntries, shard...) |
| } |
| |
| rv[partition] = combinedEntries |
| sklog.Infof("Fetched %d entries for partition %s", len(combinedEntries), partition) |
| } |
| return rv, nil |
| } |
