refactor(backend): Move credentials storage to prisma user

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/store.py

@@ -1,10 +1,10 @@
 import secrets
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from redis import Redis
-    from supabase import Client
+    from backend.executor.database import DatabaseManager
 
 from autogpt_libs.utils.synchronize import RedisKeyedMutex
 
@@ -18,8 +18,8 @@
 
 
 class SupabaseIntegrationCredentialsStore:
-    def __init__(self, supabase: "Client", redis: "Redis"):
-        self.supabase = supabase
+    def __init__(self, redis: "Redis", db: "DatabaseManager"):
+        self.db_manager: DatabaseManager = db
         self.locks = RedisKeyedMutex(redis)
 
     def add_creds(self, user_id: str, credentials: Credentials) -> None:
@@ -35,7 +35,9 @@ def add_creds(self, user_id: str, credentials: Credentials) -> None:
 
     def get_all_creds(self, user_id: str) -> list[Credentials]:
         user_metadata = self._get_user_metadata(user_id)
-        return UserMetadata.model_validate(user_metadata).integration_credentials
+        return UserMetadata.model_validate(
+            user_metadata.model_dump()
+        ).integration_credentials
 
     def get_creds_by_id(self, user_id: str, credentials_id: str) -> Credentials | None:
         all_credentials = self.get_all_creds(user_id)
@@ -90,9 +92,7 @@ def delete_creds_by_id(self, user_id: str, credentials_id: str) -> None:
             ]
             self._set_user_integration_creds(user_id, filtered_credentials)
 
-    async def store_state_token(
-        self, user_id: str, provider: str, scopes: list[str]
-    ) -> str:
+    def store_state_token(self, user_id: str, provider: str, scopes: list[str]) -> str:
         token = secrets.token_urlsafe(32)
         expires_at = datetime.now(timezone.utc) + timedelta(minutes=10)
 
@@ -105,17 +105,17 @@ async def store_state_token(
 
         with self.locked_user_metadata(user_id):
             user_metadata = self._get_user_metadata(user_id)
-            oauth_states = user_metadata.get("integration_oauth_states", [])
+            oauth_states = user_metadata.integration_oauth_states
             oauth_states.append(state.model_dump())
-            user_metadata["integration_oauth_states"] = oauth_states
+            user_metadata.integration_oauth_states = oauth_states
 
-            self.supabase.auth.admin.update_user_by_id(
-                user_id, {"user_metadata": user_metadata}
+            self.db_manager.update_user_metadata(
+                user_id=user_id, metadata=user_metadata
             )
 
         return token
 
-    async def get_any_valid_scopes_from_state_token(
+    def get_any_valid_scopes_from_state_token(
         self, user_id: str, token: str, provider: str
     ) -> list[str]:
         """
@@ -126,7 +126,7 @@ async def get_any_valid_scopes_from_state_token(
         THE CODE FOR TOKENS.
         """
         user_metadata = self._get_user_metadata(user_id)
-        oauth_states = user_metadata.get("integration_oauth_states", [])
+        oauth_states = user_metadata.integration_oauth_states
 
         now = datetime.now(timezone.utc)
         valid_state = next(
@@ -145,10 +145,10 @@ async def get_any_valid_scopes_from_state_token(
 
         return []
 
-    async def verify_state_token(self, user_id: str, token: str, provider: str) -> bool:
+    def verify_state_token(self, user_id: str, token: str, provider: str) -> bool:
         with self.locked_user_metadata(user_id):
             user_metadata = self._get_user_metadata(user_id)
-            oauth_states = user_metadata.get("integration_oauth_states", [])
+            oauth_states = user_metadata.integration_oauth_states
 
             now = datetime.now(timezone.utc)
             valid_state = next(
@@ -165,10 +165,8 @@ async def verify_state_token(self, user_id: str, token: str, provider: str) -> b
             if valid_state:
                 # Remove the used state
                 oauth_states.remove(valid_state)
-                user_metadata["integration_oauth_states"] = oauth_states
-                self.supabase.auth.admin.update_user_by_id(
-                    user_id, {"user_metadata": user_metadata}
-                )
+                user_metadata.integration_oauth_states = oauth_states
+                self.db_manager.update_user_metadata(user_id, user_metadata)
                 return True
 
         return False
@@ -177,19 +175,13 @@ def _set_user_integration_creds(
         self, user_id: str, credentials: list[Credentials]
     ) -> None:
         raw_metadata = self._get_user_metadata(user_id)
-        raw_metadata.update(
-            {"integration_credentials": [c.model_dump() for c in credentials]}
-        )
-        self.supabase.auth.admin.update_user_by_id(
-            user_id, {"user_metadata": raw_metadata}
-        )
+        raw_metadata.integration_credentials = [c.model_dump() for c in credentials]
+        self.db_manager.update_user_metadata(user_id, raw_metadata)
 
     def _get_user_metadata(self, user_id: str) -> UserMetadataRaw:
-        response = self.supabase.auth.admin.get_user_by_id(user_id)
-        if not response.user:
-            raise ValueError(f"User with ID {user_id} not found")
-        return cast(UserMetadataRaw, response.user.user_metadata)
+        metadata: UserMetadataRaw = self.db_manager.get_user_metadata(user_id=user_id)
+        return metadata
 
     def locked_user_metadata(self, user_id: str):
-        key = (self.supabase.supabase_url, f"user:{user_id}", "metadata")
+        key = (self.db_manager, f"user:{user_id}", "metadata")
         return self.locks.locked(key)

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/types.py

@@ -56,6 +56,7 @@ class OAuthState(BaseModel):
     token: str
     provider: str
     expires_at: int
+    scopes: list[str]
     """Unix timestamp (seconds) indicating when this OAuth state expires"""
 
 
@@ -64,6 +65,6 @@ class UserMetadata(BaseModel):
     integration_oauth_states: list[OAuthState] = Field(default_factory=list)
 
 
-class UserMetadataRaw(TypedDict, total=False):
-    integration_credentials: list[dict]
-    integration_oauth_states: list[dict]
+class UserMetadataRaw(BaseModel):
+    integration_credentials: list[dict] = Field(default_factory=list)
+    integration_oauth_states: list[dict] = Field(default_factory=list)

autogpt_platform/backend/Dockerfile

@@ -8,7 +8,7 @@ WORKDIR /app
 
 # Install build dependencies
 RUN apt-get update \
-    && apt-get install -y build-essential curl ffmpeg wget libcurl4-gnutls-dev libexpat1-dev gettext libz-dev libssl-dev postgresql-client git \
+    && apt-get install -y build-essential curl ffmpeg wget libcurl4-gnutls-dev libexpat1-dev libpq5 gettext libz-dev libssl-dev postgresql-client git \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 

autogpt_platform/backend/README.advanced.md

@@ -37,7 +37,7 @@ We use the Poetry to manage the dependencies. To set up the project, follow thes
 5. Generate the Prisma client
 
    ```sh
-   poetry run prisma generate --schema postgres/schema.prisma
+   poetry run prisma generate
    ```
 
 
@@ -61,7 +61,7 @@ We use the Poetry to manage the dependencies. To set up the project, follow thes
 
    ```sh
    cd ../backend
-   prisma migrate dev --schema postgres/schema.prisma
+   prisma migrate deploy
    ```
 
 ## Running The Server

autogpt_platform/backend/README.md

@@ -59,7 +59,7 @@ We use the Poetry to manage the dependencies. To set up the project, follow thes
 
    ```sh
    docker compose up db redis -d
-   poetry run prisma migrate dev 
+   poetry run prisma migrate deploy
    ```
 
 ## Running The Server

autogpt_platform/backend/backend/data/user.py

@@ -1,6 +1,8 @@
 from typing import Optional
 
+from autogpt_libs.supabase_integration_credentials_store.types import UserMetadataRaw
 from fastapi import HTTPException
+from prisma import Json
 from prisma.models import User
 
 from backend.data.db import prisma
@@ -48,3 +50,21 @@ async def create_default_user(enable_auth: str) -> Optional[User]:
             )
         return User.model_validate(user)
     return None
+
+
+async def get_user_metadata(user_id: str) -> UserMetadataRaw:
+    user = await User.prisma().find_unique_or_raise(
+        where={"id": user_id},
+    )
+    return (
+        UserMetadataRaw.model_validate(user.metadata)
+        if user.metadata
+        else UserMetadataRaw()
+    )
+
+
+async def update_user_metadata(user_id: str, metadata: UserMetadataRaw):
+    await User.prisma().update(
+        where={"id": user_id},
+        data={"metadata": Json(metadata.model_dump())},
+    )

autogpt_platform/backend/backend/executor/database.py

@@ -16,6 +16,7 @@
 )
 from backend.data.graph import get_graph, get_node
 from backend.data.queue import RedisEventQueue
+from backend.data.user import get_user_metadata, update_user_metadata
 from backend.util.service import AppService, expose
 from backend.util.settings import Config
 
@@ -73,3 +74,7 @@ def wrapper(self, *args: P.args, **kwargs: P.kwargs) -> R:
         Callable[[Any, str, int, str, dict[str, str], float, float], int],
         exposed_run_and_wait(user_credit_model.spend_credits),
     )
+
+    # User + User Metadata
+    get_user_metadata = exposed_run_and_wait(get_user_metadata)
+    update_user_metadata = exposed_run_and_wait(update_user_metadata)

autogpt_platform/backend/backend/executor/manager.py

@@ -31,7 +31,7 @@
 from backend.data.model import CREDENTIALS_FIELD_NAME, CredentialsMetaInput
 from backend.integrations.creds_manager import IntegrationCredentialsManager
 from backend.util import json
-from backend.util.cache import thread_cached_property
+from backend.util.cache import thread_cached
 from backend.util.decorator import error_logged, time_measured
 from backend.util.logging import configure_logging
 from backend.util.process import set_service_name
@@ -417,7 +417,7 @@ def on_node_executor_start(cls):
         redis.connect()
         cls.pid = os.getpid()
         cls.db_client = get_db_client()
-        cls.creds_manager = IntegrationCredentialsManager()
+        cls.creds_manager = IntegrationCredentialsManager(db_manager=cls.db_client)
 
         # Set up shutdown handlers
         cls.shutdown_lock = threading.Lock()
@@ -670,7 +670,7 @@ def run_service(self):
         )
 
         self.credentials_store = SupabaseIntegrationCredentialsStore(
-            self.supabase, redis.get_redis()
+            redis=redis.get_redis(), db=self.db_client
         )
         self.executor = ProcessPoolExecutor(
             max_workers=self.pool_size,
@@ -701,7 +701,7 @@ def cleanup(self):
 
         super().cleanup()
 
-    @thread_cached_property
+    @property
     def db_client(self) -> "DatabaseManager":
         return get_db_client()
 
@@ -857,6 +857,7 @@ def _validate_node_input_credentials(self, graph: Graph, user_id: str):
 # ------- UTILITIES ------- #
 
 
+@thread_cached
 def get_db_client() -> "DatabaseManager":
     from backend.executor import DatabaseManager
 

autogpt_platform/backend/backend/integrations/creds_manager.py

@@ -10,11 +10,10 @@
 from redis.lock import Lock as RedisLock
 
 from backend.data import redis
+from backend.executor.database import DatabaseManager
 from backend.integrations.oauth import HANDLERS_BY_NAME, BaseOAuthHandler
 from backend.util.settings import Settings
 
-from ..server.integrations.utils import get_supabase
-
 logger = logging.getLogger(__name__)
 settings = Settings()
 
@@ -51,10 +50,12 @@ class IntegrationCredentialsManager:
     cause so much latency that it's worth implementing.
     """
 
-    def __init__(self):
+    def __init__(self, db_manager: DatabaseManager):
         redis_conn = redis.get_redis()
         self._locks = RedisKeyedMutex(redis_conn)
-        self.store = SupabaseIntegrationCredentialsStore(get_supabase(), redis_conn)
+        self.store = SupabaseIntegrationCredentialsStore(
+            redis=redis_conn, db=db_manager
+        )
 
     def create(self, user_id: str, credentials: Credentials) -> None:
         return self.store.add_creds(user_id, credentials)
@@ -131,7 +132,7 @@ def delete(self, user_id: str, credentials_id: str) -> None:
 
     def _acquire_lock(self, user_id: str, credentials_id: str, *args: str) -> RedisLock:
         key = (
-            self.store.supabase.supabase_url,
+            self.store.db_manager,
             f"user:{user_id}",
             f"credentials:{credentials_id}",
             *args,

autogpt_platform/backend/backend/server/integrations/router.py

@@ -10,6 +10,7 @@
 from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request
 from pydantic import BaseModel, Field, SecretStr
 
+from backend.executor.manager import get_db_client
 from backend.integrations.creds_manager import IntegrationCredentialsManager
 from backend.integrations.oauth import HANDLERS_BY_NAME, BaseOAuthHandler
 from backend.util.settings import Settings
@@ -19,7 +20,8 @@
 logger = logging.getLogger(__name__)
 settings = Settings()
 router = APIRouter()
-creds_manager = IntegrationCredentialsManager()
+
+creds_manager = IntegrationCredentialsManager(db_manager=get_db_client())
 
 
 class LoginResponse(BaseModel):
@@ -41,7 +43,7 @@ async def login(
     requested_scopes = scopes.split(",") if scopes else []
 
     # Generate and store a secure random state token along with the scopes
-    state_token = await creds_manager.store.store_state_token(
+    state_token = creds_manager.store.store_state_token(
         user_id, provider, requested_scopes
     )
 
@@ -70,12 +72,12 @@ async def callback(
     handler = _get_provider_oauth_handler(request, provider)
 
     # Verify the state token
-    if not await creds_manager.store.verify_state_token(user_id, state_token, provider):
+    if not creds_manager.store.verify_state_token(user_id, state_token, provider):
         logger.warning(f"Invalid or expired state token for user {user_id}")
         raise HTTPException(status_code=400, detail="Invalid or expired state token")
 
     try:
-        scopes = await creds_manager.store.get_any_valid_scopes_from_state_token(
+        scopes = creds_manager.store.get_any_valid_scopes_from_state_token(
             user_id, state_token, provider
         )
         logger.debug(f"Retrieved scopes from state token: {scopes}")

autogpt_platform/backend/backend/server/rest_api.py

@@ -19,6 +19,7 @@
 from backend.data.credit import get_block_costs, get_user_credit_model
 from backend.data.user import get_or_create_user
 from backend.executor import ExecutionManager, ExecutionScheduler
+from backend.executor.manager import get_db_client
 from backend.integrations.creds_manager import IntegrationCredentialsManager
 from backend.server.model import CreateGraph, SetGraphActiveVersion
 from backend.util.cache import thread_cached_property
@@ -97,7 +98,7 @@ def run_service(self):
             tags=["integrations"],
             dependencies=[Depends(auth_middleware)],
         )
-        self.integration_creds_manager = IntegrationCredentialsManager()
+        self.integration_creds_manager = IntegrationCredentialsManager(get_db_client())
 
         api_router.include_router(
             backend.server.routers.analytics.router,