From 3a636c0db0099a2b69964e404d97b40ec41ea96e Mon Sep 17 00:00:00 2001
From: Boris Arzentar <borisarzentar@gmail.com>
Date: Wed, 16 Oct 2024 22:33:04 +0200
Subject: [PATCH] fix: add default user via migration

---
 .../versions/482cd6517ce4_add_default_user.py | 27 +++++++++++++++++++
 cognee/modules/users/methods/__init__.py      |  1 +
 cognee/modules/users/methods/create_user.py   |  1 -
 cognee/modules/users/methods/delete_user.py   | 17 ++++++++++++
 4 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 alembic/versions/482cd6517ce4_add_default_user.py
 create mode 100644 cognee/modules/users/methods/delete_user.py

diff --git a/alembic/versions/482cd6517ce4_add_default_user.py b/alembic/versions/482cd6517ce4_add_default_user.py
new file mode 100644
index 000000000..ec916829e
--- /dev/null
+++ b/alembic/versions/482cd6517ce4_add_default_user.py
@@ -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"))
diff --git a/cognee/modules/users/methods/__init__.py b/cognee/modules/users/methods/__init__.py
index 1f8127a4c..80b2a8a2a 100644
--- a/cognee/modules/users/methods/__init__.py
+++ b/cognee/modules/users/methods/__init__.py
@@ -1,5 +1,6 @@
 from .get_user import get_user
 from .create_user import create_user
+from .delete_user import delete_user
 from .get_default_user import get_default_user
 from .create_default_user import create_default_user
 from .get_authenticated_user import get_authenticated_user
diff --git a/cognee/modules/users/methods/create_user.py b/cognee/modules/users/methods/create_user.py
index 3145fd799..cfe1dbd03 100644
--- a/cognee/modules/users/methods/create_user.py
+++ b/cognee/modules/users/methods/create_user.py
@@ -32,7 +32,6 @@ async def create_user(
                         await session.refresh(user)
 
                     return user
-                    print(f"User created: {user.email}")
     except UserAlreadyExists as error:
         print(f"User {email} already exists")
         raise error
diff --git a/cognee/modules/users/methods/delete_user.py b/cognee/modules/users/methods/delete_user.py
new file mode 100644
index 000000000..dd0a585b6
--- /dev/null
+++ b/cognee/modules/users/methods/delete_user.py
@@ -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