Fix migration (#337)

MinBZK · Oct 30, 2024 · f3eb789 · f3eb789
2 parents d3b51d4 + bc4d3d3
commit f3eb789
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git 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
@@ -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: