From ca1cf242dcb2893962fe529569a301ef8234e343 Mon Sep 17 00:00:00 2001
From: maxdml <maxdml@dbos.dev>
Date: Tue, 2 Sep 2025 09:59:55 -0700
Subject: [PATCH] handle more drivers

---
 dbos/system_database.go | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/dbos/system_database.go b/dbos/system_database.go
index 7c27d6a3..14066d63 100644
--- a/dbos/system_database.go
+++ b/dbos/system_database.go
@@ -140,21 +140,22 @@ const (
 
 func runMigrations(databaseURL string) error {
 	// Change the driver to pgx5
-	databaseURL = "pgx5://" + strings.TrimPrefix(databaseURL, "postgres://")
-
-	// Create migration source from embedded files
-	d, err := iofs.New(migrationFiles, "migrations")
-	if err != nil {
-		return newInitializationError(fmt.Sprintf("failed to create migration source: %v", err))
-	}
-
-	// Add custom migration table name to avoid conflicts with user migrations
-	// Parse the URL to properly determine where to add the query parameter
 	parsedURL, err := url.Parse(databaseURL)
 	if err != nil {
 		return newInitializationError(fmt.Sprintf("failed to parse database URL: %v", err))
 	}
+	// Handle various PostgreSQL URL schemes
+	switch parsedURL.Scheme {
+	case "postgres", "postgresql":
+		parsedURL.Scheme = "pgx5"
+	case "pgx5":
+		// Already in correct format
+	default:
+		return newInitializationError(fmt.Sprintf("unsupported database URL scheme: %s", parsedURL.Scheme))
+	}
+	databaseURL = parsedURL.String()
 
+	// Add custom migration table name to avoid conflicts with user migrations
 	// Check if query parameters already exist
 	separator := "?"
 	if parsedURL.RawQuery != "" {
@@ -162,6 +163,12 @@ func runMigrations(databaseURL string) error {
 	}
 	databaseURL += separator + "x-migrations-table=" + _DBOS_MIGRATION_TABLE
 
+	// Create migration source from embedded files
+	d, err := iofs.New(migrationFiles, "migrations")
+	if err != nil {
+		return newInitializationError(fmt.Sprintf("failed to create migration source: %v", err))
+	}
+
 	// Create migrator
 	m, err := migrate.NewWithSourceInstance("iofs", d, databaseURL)
 	if err != nil {