fix: minor bugs that prevented connecting and migrating to dbs and auth

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/store.py

@@ -23,12 +23,12 @@
 
 class SupabaseIntegrationCredentialsStore:
     def __init__(self, redis: "Redis"):
-        self.prisma: Prisma = db.prisma
+        self.prisma: Prisma = Prisma()
         self.locks = RedisKeyedMutex(redis)
 
     async def add_creds(self, user_id: str, credentials: Credentials) -> None:
         with self.locked_user_metadata(user_id):
-            if self.get_creds_by_id(user_id, credentials.id):
+            if await self.get_creds_by_id(user_id, credentials.id):
                 raise ValueError(
                     f"Can not re-create existing credentials #{credentials.id} "
                     f"for user #{user_id}"
@@ -43,7 +43,9 @@ async def get_all_creds(self, user_id: str) -> list[Credentials]:
             user_metadata.model_dump()
         ).integration_credentials
 
-    async def get_creds_by_id(self, user_id: str, credentials_id: str) -> Credentials | None:
+    async def get_creds_by_id(
+        self, user_id: str, credentials_id: str
+    ) -> Credentials | None:
         all_credentials = await self.get_all_creds(user_id)
         return next((c for c in all_credentials if c.id == credentials_id), None)
 
@@ -117,8 +119,11 @@ async def store_state_token(
             oauth_states.append(state.model_dump())
             user_metadata.integration_oauth_states = oauth_states
 
+            if not self.prisma.is_connected():
+                await self.prisma.connect()
             await self.prisma.user.update(
-                where={"id": user_id}, data={"metadata": Json(user_metadata.model_dump())}
+                where={"id": user_id},
+                data={"metadata": Json(user_metadata.model_dump())},
             )
 
         return token
@@ -174,6 +179,8 @@ async def verify_state_token(self, user_id: str, token: str, provider: str) -> b
                 # Remove the used state
                 oauth_states.remove(valid_state)
                 user_metadata.integration_oauth_states = oauth_states
+                if not self.prisma.is_connected():
+                    await self.prisma.connect()
                 await self.prisma.user.update(
                     where={"id": user_id},
                     data={"metadata": Json(user_metadata.model_dump())},
@@ -187,11 +194,15 @@ async def _set_user_integration_creds(
     ) -> None:
         raw_metadata = await self._get_user_metadata(user_id)
         raw_metadata.integration_credentials = [c.model_dump() for c in credentials]
+        if not self.prisma.is_connected():
+            await self.prisma.connect()
         await self.prisma.user.update(
             where={"id": user_id}, data={"metadata": Json(raw_metadata.model_dump())}
         )
 
     async def _get_user_metadata(self, user_id: str) -> UserMetadataRaw:
+        if not self.prisma.is_connected():
+            await self.prisma.connect()
         user = await self.prisma.user.find_unique(where={"id": user_id})
         if not user:
             raise ValueError(f"User with ID {user_id} not found")
@@ -202,5 +213,5 @@ async def _get_user_metadata(self, user_id: str) -> UserMetadataRaw:
         )
 
     def locked_user_metadata(self, user_id: str):
-        key = ("usermetadatalock", f"user:{user_id}", "metadata")
+        key = (self.prisma, f"user:{user_id}", "metadata")
         return self.locks.locked(key)

autogpt_platform/backend/backend/integrations/creds_manager.py

@@ -130,7 +130,7 @@ async def delete(self, user_id: str, credentials_id: str) -> None:
 
     def _acquire_lock(self, user_id: str, credentials_id: str, *args: str) -> RedisLock:
         key = (
-            "usermetadatalock",
+            self.store.prisma,
             f"user:{user_id}",
             f"credentials:{credentials_id}",
             *args,

autogpt_platform/backend/backend/util/settings.py

@@ -153,6 +153,11 @@ class Config(UpdateTrackingModel["Config"], BaseSettings):
         description="What environment to behave as: local or cloud",
     )
 
+    direct_database_url: str = Field(
+        default="",
+        description="The URL for the direct database. This is used to run migrations.",
+    )
+
     backend_cors_allow_origins: List[str] = Field(default_factory=list)
 
     @field_validator("backend_cors_allow_origins")

autogpt_platform/backend/migrate/run_migrations.py

@@ -2,6 +2,7 @@
 import os
 import subprocess
 
+from backend.util.settings import Settings
 import psycopg2
 
 
@@ -17,7 +18,11 @@ def run_prisma_migrations():
 
 
 def run_custom_migrations():
-    db_url = os.environ.get("DIRECT_DATABASE_URL")
+    settings = Settings()
+    db_url = settings.config.direct_database_url
+    if not db_url:
+        logging.error("No Direct DB URL Provided")
+        return
     conn = psycopg2.connect(db_url)
 
     try:

autogpt_platform/backend/pyproject.toml

@@ -70,6 +70,7 @@ cli = "backend.cli:main"
 format = "linter:format"
 lint = "linter:lint"
 test = "run_tests:test"
+migrate = "migrate.run_migrations:main"
 migrate_db = "migrate.run_migrations:main"
 
 [tool.pytest-watcher]