fix: add default user via migration

This commit is contained in:
Boris Arzentar 2024-10-16 22:33:04 +02:00
parent 2485eb76fa
commit 3a636c0db0
4 changed files with 45 additions and 1 deletions

View file

@ -0,0 +1,27 @@
"""Add default user
Revision ID: 482cd6517ce4
Revises: 8057ae7329c2
Create Date: 2024-10-16 22:17:18.634638
"""
from typing import Sequence, Union
from sqlalchemy.util import await_only
from cognee.modules.users.methods import create_default_user, delete_user
# revision identifiers, used by Alembic.
revision: str = '482cd6517ce4'
down_revision: Union[str, None] = '8057ae7329c2'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
await_only(create_default_user())
def downgrade() -> None:
await_only(delete_user("default_user@example.com"))

View file

@ -1,5 +1,6 @@
from .get_user import get_user from .get_user import get_user
from .create_user import create_user from .create_user import create_user
from .delete_user import delete_user
from .get_default_user import get_default_user from .get_default_user import get_default_user
from .create_default_user import create_default_user from .create_default_user import create_default_user
from .get_authenticated_user import get_authenticated_user from .get_authenticated_user import get_authenticated_user

View file

@ -32,7 +32,6 @@ async def create_user(
await session.refresh(user) await session.refresh(user)
return user return user
print(f"User created: {user.email}")
except UserAlreadyExists as error: except UserAlreadyExists as error:
print(f"User {email} already exists") print(f"User {email} already exists")
raise error raise error

View file

@ -0,0 +1,17 @@
from fastapi_users.exceptions import UserNotExists
from cognee.infrastructure.databases.relational import get_relational_engine
from ..get_user_manager import get_user_manager_context
from ..get_user_db import get_user_db_context
async def delete_user(email: str):
try:
relational_engine = get_relational_engine()
async with relational_engine.get_async_session() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
user = await user_manager.get_by_email(email)
await user_manager.delete(user)
except UserNotExists as error:
print(f"User {email} doesn't exist")
raise error