From bc4d3d3e5b03a8597d8f922b76925fbb43ccb4d0 Mon Sep 17 00:00:00 2001
From: Berry den Hartog <38954346+berrydenhartog@users.noreply.github.com>
Date: Tue, 29 Oct 2024 22:12:03 +0100
Subject: [PATCH] Fix migration

---
 amt/core/config.py    |  6 ++++--
 amt/migrations/env.py | 19 ++++++++++++++-----
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/amt/core/config.py b/amt/core/config.py
index 6a8413bc..00c72a16 100644
--- a/amt/core/config.py
+++ b/amt/core/config.py
@@ -44,7 +44,7 @@ class Settings(BaseSettings):
 
     # todo(berry): create submodel for database settings
     APP_DATABASE_SCHEME: DatabaseSchemaType = "sqlite"
-    APP_DATABASE_DRIVER: str | None = "aiosqlite"
+    APP_DATABASE_DRIVER: str | None = None
 
     APP_DATABASE_SERVER: str = "db"
     APP_DATABASE_PORT: int = 5432
@@ -68,10 +68,12 @@ def SQLALCHEMY_ECHO(self) -> bool:
 
     @computed_field
     def SQLALCHEMY_DATABASE_URI(self) -> str:
+        default_driver: str = "aiosqlite" if self.APP_DATABASE_SCHEME == "sqlite" else "asyncpg"
+
         scheme: str = (
             f"{self.APP_DATABASE_SCHEME}+{self.APP_DATABASE_DRIVER}"
             if isinstance(self.APP_DATABASE_DRIVER, str)
-            else self.APP_DATABASE_SCHEME
+            else f"{self.APP_DATABASE_SCHEME}+{default_driver}"
         )
 
         if self.APP_DATABASE_SCHEME == "sqlite":
diff --git a/amt/migrations/env.py b/amt/migrations/env.py
index 9f4b92b0..53309e78 100644
--- a/amt/migrations/env.py
+++ b/amt/migrations/env.py
@@ -19,19 +19,28 @@
 
 
 def get_url() -> str:
-    scheme = os.getenv("APP_DATABASE_SCHEME", "sqlite")
-    driver = os.getenv("APP_DATABASE_DRIVER", "aiosqlite")
+    scheme_temp = os.getenv("APP_DATABASE_SCHEME", "sqlite")
+    driver = os.getenv("APP_DATABASE_DRIVER", None)
+    default_driver = "aiosqlite" if scheme_temp == "sqlite" else "asyncpg"
+
+    scheme = (
+            f"{scheme_temp}+{driver}"
+            if isinstance(driver, str)
+            else f"{scheme_temp}+{default_driver}"
+    )
+
+    if scheme_temp == "sqlite":
 
-    if scheme == "sqlite":
         file = os.getenv("APP_DATABASE_FILE", "database.sqlite3")
-        return f"{scheme}+{driver}:///{file}"
+        return f"{scheme}:///{file}"
+
 
     user = os.getenv("APP_DATABASE_USER", "amt")
     password = os.getenv("APP_DATABASE_PASSWORD", "")
     server = os.getenv("APP_DATABASE_SERVER", "db")
     port = os.getenv("APP_DATABASE_PORT", "5432")
     db = os.getenv("APP_DATABASE_DB", "amt")
-    return f"{scheme}+{driver}://{user}:{password}@{server}:{port}/{db}"
+    return f"{scheme}://{user}:{password}@{server}:{port}/{db}"
 
 
 def run_migrations_offline() -> None: