Skip to content

Commit

Permalink
Revert "fix(platform): Remove migrate and encrypt function" (#8654)
Browse files Browse the repository at this point in the history
Reverts c707ee9 (#8646)

The problem analysis that led to #8646 contained some errors, so the migration removed in the PR doesn't seem to have been the cause of the problem we were hunting. Also, this migration is an essential part of the security improvement that we made 2 weeks ago.
  • Loading branch information
Pwuts authored Nov 14, 2024
1 parent e140873 commit f9633ff
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 0 deletions.
40 changes: 40 additions & 0 deletions autogpt_platform/backend/backend/data/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,3 +94,43 @@ async def update_user_integrations(user_id: str, data: UserIntegrations):
where={"id": user_id},
data={"integrations": encrypted_data},
)


async def migrate_and_encrypt_user_integrations():
"""Migrate integration credentials and OAuth states from metadata to integrations column."""
users = await User.prisma().find_many(
where={
"metadata": {
"path": ["integration_credentials"],
"not": Json({"a": "yolo"}), # bogus value works to check if key exists
} # type: ignore
}
)
logger.info(f"Migrating integration credentials for {len(users)} users")

for user in users:
raw_metadata = cast(UserMetadataRaw, user.metadata)
metadata = UserMetadata.model_validate(raw_metadata)

# Get existing integrations data
integrations = await get_user_integrations(user_id=user.id)

# Copy credentials and oauth states from metadata if they exist
if metadata.integration_credentials and not integrations.credentials:
integrations.credentials = metadata.integration_credentials
if metadata.integration_oauth_states:
integrations.oauth_states = metadata.integration_oauth_states

# Save to integrations column
await update_user_integrations(user_id=user.id, data=integrations)

# Remove from metadata
raw_metadata = dict(raw_metadata)
raw_metadata.pop("integration_credentials", None)
raw_metadata.pop("integration_oauth_states", None)

# Update metadata without integration data
await User.prisma().update(
where={"id": user.id},
data={"metadata": Json(raw_metadata)},
)
1 change: 1 addition & 0 deletions autogpt_platform/backend/backend/server/rest_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
async def lifespan_context(app: fastapi.FastAPI):
await backend.data.db.connect()
await backend.data.block.initialize_blocks()
await backend.data.user.migrate_and_encrypt_user_integrations()
yield
await backend.data.db.disconnect()

Expand Down

0 comments on commit f9633ff

Please sign in to comment.