Cast ids to misk.hibernate.Id instead of Long in BulkShardMigrator

cashapp · Aug 7, 2019 · 89b1861 · 89b1861
1 parent 4ee2b86
commit 89b1861
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/misk-hibernate/src/main/kotlin/misk/hibernate/migrate/BulkShardMigrator.kt b/misk-hibernate/src/main/kotlin/misk/hibernate/migrate/BulkShardMigrator.kt
@@ -206,10 +206,9 @@ class BulkShardMigrator<R : DbRoot<R>, C : DbChild<R, C>> private constructor(
             "SELECT id FROM $tableName WHERE $rootColumnName = :target AND id IN (:ids)"
         ).setParameter("target", targetRoot!!.id)
           .setParameterList("ids", sourceRecords.keys)
-          .list()
-          .stream()
-          .map { it as Long }
-          .collect(Collectors.toSet<Long>())
+            .resultList
+            .map { (it as Id<*>).id }
+            .toSet()
 
         logger.info(
             "Bulk migrating (distinct shard) ${sourceRecords.size} entities for table $tableName"