From 52e3087bb45acbe55458e9bd82adcd1504dde76a Mon Sep 17 00:00:00 2001
From: Adrian Ancona Novelo <soonick5@yahoo.com.mx>
Date: Sat, 16 Sep 2017 14:24:22 +0800
Subject: [PATCH] Only import the correct driver when doing a migration. Fixes
 #447

---
 cmd/commands/migrate/migrate.go | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/cmd/commands/migrate/migrate.go b/cmd/commands/migrate/migrate.go
index 1870945..92fa16e 100644
--- a/cmd/commands/migrate/migrate.go
+++ b/cmd/commands/migrate/migrate.go
@@ -196,6 +196,17 @@ func checkForSchemaUpdateTable(db *sql.DB, driver string) {
 	}
 }
 
+func driverImportStatement(driver string) string {
+	switch driver {
+	case "mysql":
+		return "github.com/go-sql-driver/mysql"
+	case "postgres":
+		return "github.com/lib/pq"
+	default:
+		return "github.com/go-sql-driver/mysql"
+	}
+}
+
 func showMigrationsTableSQL(driver string) string {
 	switch driver {
 	case "mysql":
@@ -263,6 +274,7 @@ func writeMigrationSourceFile(dir, source, driver, connStr string, latestTime in
 		beeLogger.Log.Fatalf("Could not create file: %s", err)
 	} else {
 		content := strings.Replace(MigrationMainTPL, "{{DBDriver}}", driver, -1)
+		content = strings.Replace(content, "{{DriverRepo}}", driverImportStatement(driver), -1)
 		content = strings.Replace(content, "{{ConnStr}}", connStr, -1)
 		content = strings.Replace(content, "{{LatestTime}}", strconv.FormatInt(latestTime, 10), -1)
 		content = strings.Replace(content, "{{LatestName}}", latestName, -1)
@@ -346,8 +358,7 @@ import(
 	"github.com/astaxie/beego/orm"
 	"github.com/astaxie/beego/migration"
 
-	_ "github.com/go-sql-driver/mysql"
-	_ "github.com/lib/pq"
+	_ "{{DriverRepo}}"
 )
 
 func init(){