feat: add down migration for UUID mapping

internal/persistence/sql/migrations/migratest/migration_test.go

@@ -150,9 +150,6 @@ func TestMigrations(t *testing.T) {
 			// cleanup first
 			require.NoError(t, tm.Down(ctx, -1))
 
-			t.Log("before migration")
-			logMigrationStatus(t, tm)
-
 			t.Run("suite=up", func(t *testing.T) {
 				if err := tm.Up(ctx); err != nil {
 					t.Log("migrations failed:", err)
@@ -161,12 +158,15 @@ func TestMigrations(t *testing.T) {
 				logMigrationStatus(t, tm)
 			})
 
+			reg := driver.NewTestRegistry(t, db)
+			require.NoError(t,
+				reg.Config(ctx).Set(config.KeyNamespaces, []*namespace.Namespace{
+					{ID: 1, Name: "foo"},
+					{ID: 2, Name: "uuid_test"},
+				}))
+			p, err := sql.NewPersister(ctx, reg, uuid.Must(uuid.FromString("77fdc5e0-2260-49da-8aae-c36ba255d05b")))
+
 			t.Run("suite=fixtures", func(t *testing.T) {
-				reg := driver.NewTestRegistry(t, db)
-				require.NoError(t,
-					reg.Config(ctx).Set(config.KeyNamespaces, []*namespace.Namespace{
-						{ID: 1, Name: "foo"}}))
-				p, err := sql.NewPersister(ctx, reg, uuid.Must(uuid.FromString("77fdc5e0-2260-49da-8aae-c36ba255d05b")))
 
 				t.Run("table=legacy namespaces", func(t *testing.T) {
 					// as they are legacy, we expect them to be actually dropped
@@ -211,6 +211,32 @@ func TestMigrations(t *testing.T) {
 				})
 			})
 
+			t.Run("suite=uuid_migrations", func(t *testing.T) {
+				require.NoError(t, tm.Down(ctx, -1))
+
+				// Migrate up to (including) "migrate-strings-to-uuids"
+				migrateTo(t, tm, "migrate-strings-to-uuids")
+				t.Log("status after up migration")
+				logMigrationStatus(t, tm)
+
+				// Assert that relationtuples have UUIDs
+				tuples, _, err := p.GetRelationTuples(ctx, &relationtuple.RelationQuery{Namespace: "uuid_test"})
+				require.NoError(t, err)
+				assertIsUUID(t, *tuples[0].Subject.SubjectID())
+				assertIsUUID(t, tuples[0].Object)
+
+				// Migrate down to before "migrate-strings-to-uuids"
+				require.NoError(t, tm.Down(ctx, 1))
+				t.Log("status after down migration")
+				logMigrationStatus(t, tm)
+
+				// Assert that relationtuples have strings
+				tuples, _, err = p.GetRelationTuples(ctx, &relationtuple.RelationQuery{Namespace: "uuid_test"})
+				require.NoError(t, err)
+				assert.Equal(t, "user", *tuples[0].Subject.SubjectID())
+				assert.Equal(t, "object", tuples[0].Object)
+			})
+
 			t.Run("suite=down", func(t *testing.T) {
 				if debugOnDisk && db.Name == "sqlite" {
 					t.SkipNow()
@@ -221,7 +247,30 @@ func TestMigrations(t *testing.T) {
 	}
 }
 
+func migrateTo(t *testing.T, tm *popx.MigrationBox, name string) {
+	statuses, err := tm.Status(context.Background())
+	require.NoError(t, err)
+
+	for i, status := range statuses {
+		if status.Name == name {
+			_, err = tm.UpTo(context.Background(), i+1)
+			require.NoError(t, err)
+			return
+		}
+	}
+	t.Fatal("could not find ", name)
+}
+
+func assertIsUUID(t *testing.T, id string) {
+	t.Helper()
+	if _, err := uuid.FromString(id); err != nil {
+		t.Errorf("expected %q to be a UUID", id)
+	}
+}
+
 func logMigrationStatus(t *testing.T, m *popx.MigrationBox) {
+	t.Helper()
+
 	status, err := m.Status(context.Background())
 	require.NoError(t, err)
 	s := strings.Builder{}

internal/persistence/sql/migrations/migratest/testdata/20210623162417_testdata.sql

@@ -7,3 +7,8 @@ INSERT INTO keto_relation_tuples (shard_id, nid, namespace_id, object, relation,
                                   subject_set_object, subject_set_relation, commit_time)
 VALUES ('1457d171-4473-4c77-bd88-0809529c599f', '77fdc5e0-2260-49da-8aae-c36ba255d05b', 1, 'object', 'relation', NULL,
         1, 's_object', 's_relation', '2021-09-30 15:50:52');
+
+INSERT INTO keto_relation_tuples (shard_id, nid, namespace_id, object, relation, subject_id, subject_set_namespace_id,
+                                  subject_set_object, subject_set_relation, commit_time)
+VALUES ('7a9d6df4-3b60-4cae-996e-d15d78b0fc35', '77fdc5e0-2260-49da-8aae-c36ba255d05b', 2, 'object', 'relation', 'user',
+        NULL, NULL, NULL, '2021-09-30 15:50:52');

internal/persistence/sql/migrations/uuidmapping/uuid_mapping_migrator.go

@@ -43,6 +43,8 @@ var (
 	name       = "migrate-strings-to-uuids"
 	version    = "20220513210000000000"
 	Migrations = popx.Migrations{
+		// The "up" migration will add the UUID mappings to the database and
+		// replace the strings with UUIDs.
 		{
 			Version:   version,
 			Name:      name,
@@ -58,6 +60,7 @@ var (
 					}
 
 					for _, rt := range relationTuples {
+						rt := rt
 						if err = migrateSubjectID(conn, &rt); err != nil {
 							return fmt.Errorf("could not migrate subject ID: %w", err)
 						}
@@ -80,16 +83,44 @@ var (
 				return nil
 			},
 		},
-		// We have to specify the "down" migration even if it is a no-op, since
-		// the migrator will still manipulate the version table in the database.
+		// The "down" migration will replace all UUIDs with strings from the
+		// mapping table.
 		{
 			Version:   version,
 			Name:      name,
 			Path:      name,
 			Direction: "down",
 			DBType:    "all",
 			Type:      "go",
-			Runner: func(_ popx.Migration, _ *pop.Connection, _ *pop.Tx) error {
+			Runner: func(_ popx.Migration, conn *pop.Connection, _ *pop.Tx) error {
+				for page := 1; ; page++ {
+					relationTuples, hasNext, err := getRelationTuples(conn, page)
+					if err != nil {
+						return fmt.Errorf("could not get relation tuples: %w", err)
+					}
+
+					for _, rt := range relationTuples {
+						rt := rt
+						fields := []*string{&rt.Object}
+						if rt.SubjectID.Valid {
+							fields = append(fields, &rt.SubjectID.String)
+						}
+						if rt.SubjectSetObject.Valid {
+							fields = append(fields, &rt.SubjectSetObject.String)
+						}
+						if err := batchReplaceUUIDs(conn, fields); err != nil {
+							return fmt.Errorf("could not replace UUIDs: %w", err)
+						}
+						if err = conn.Update(&rt); err != nil {
+							return fmt.Errorf("failed to update relation tuple: %w", err)
+						}
+					}
+
+					if !hasNext {
+						break
+					}
+				}
+
 				return nil
 			},
 		},
@@ -188,3 +219,50 @@ func getRelationTuples(conn *pop.Connection, page int) (
 	}
 	return res, q.Paginator.Page < q.Paginator.TotalPages, nil
 }
+
+func removeNonUUIDs(fields []*string) []*string {
+	var res []*string
+	for _, f := range fields {
+		if f == nil || *f == "" {
+			continue
+		}
+		if _, err := uuid.FromString(*f); err != nil {
+			continue
+		}
+		res = append(res, f)
+	}
+	return res
+}
+
+func batchReplaceUUIDs(conn *pop.Connection, ids []*string) (err error) {
+	ids = removeNonUUIDs(ids)
+
+	if len(ids) == 0 {
+		return
+	}
+
+	// Build a map from UUID -> target
+	uuidToTargets := make(map[string][]*string)
+	for _, id := range ids {
+		if ids, ok := uuidToTargets[*id]; ok {
+			uuidToTargets[*id] = append(ids, id)
+		} else {
+			uuidToTargets[*id] = []*string{id}
+		}
+	}
+
+	mappings := &[]UUIDMapping{}
+	query := conn.Where("id in (?)", ids)
+	if err := sqlcon.HandleError(query.All(mappings)); err != nil {
+		return err
+	}
+
+	// Write the representation to the correct index.
+	for _, m := range *mappings {
+		for _, target := range uuidToTargets[m.ID.String()] {
+			*target = m.StringRepresentation
+		}
+	}
+
+	return
+}

internal/persistence/sql/uuid_mapping_test.go

@@ -22,6 +22,7 @@ func assertCheckErr(t assert.TestingT, err error, msgAndArgs ...interface{}) boo
 	}
 
 	if strings.Contains(err.Error(), "keto_uuid_mappings") || // <- normal databases
+		strings.Contains(err.Error(), "SQLSTATE 23505") || // <- cockroach
 		strings.Contains(err.Error(), "SQLSTATE 23514") { // <- mysql
 		return true
 	}