diff --git a/vertx-demo/src/test/kotlin/app/GenerateMigration.kt b/vertx-demo/src/test/kotlin/app/GenerateMigration.kt
index 93d7cb8..cacb6a1 100644
--- a/vertx-demo/src/test/kotlin/app/GenerateMigration.kt
+++ b/vertx-demo/src/test/kotlin/app/GenerateMigration.kt
@@ -66,11 +66,8 @@ object GenerateMigration {
     // 创建注解映射器
     val mapper = createMapper()
 
-    // 扫描实体类
-    val entityClasses = scanEntityClasses(entityPackage)
-
     // 生成迁移文件
-    SqlMigrationGenerator.generateMigrations(entityClasses, mapper)
+    SqlMigrationGenerator.generateMigrations(entityPackage, mapper)
 
     println("数据库迁移文件生成完成")
   }
diff --git a/vertx-fw/src/main/kotlin/org/aikrai/vertx/db/migration/SqlMigrationGenerator.kt b/vertx-fw/src/main/kotlin/org/aikrai/vertx/db/migration/SqlMigrationGenerator.kt
index a1d276f..775fe90 100644
--- a/vertx-fw/src/main/kotlin/org/aikrai/vertx/db/migration/SqlMigrationGenerator.kt
+++ b/vertx-fw/src/main/kotlin/org/aikrai/vertx/db/migration/SqlMigrationGenerator.kt
@@ -34,7 +34,9 @@ class SqlMigrationGenerator {
          * @param entityClasses 实体类列表
          * @param mapper 注解映射中间类
          */
-        fun generateMigrations(entityClasses: List<KClass<*>>, mapper: SqlAnnotationMapper) {
+        fun generateMigrations(entityPackage: String, mapper: SqlAnnotationMapper) {
+            val entityClasses = scanEntityClasses(entityPackage)
+
             // 创建必要的目录
             createDirectories()