Skip to content

Commit 63f6a5e

Browse files
fix(management,createsuperuser): fix ORM model registration and async/sync operations
- Add automatic model import in createsuperuser command - Fix async/sync operation conflicts by forcing sync mode - Ensure proper model registration in ModelMeta before ORM operations - Fix group creation and user creation through ORM - Improve error handling for duplicate group creation
1 parent 9b57cd2 commit 63f6a5e

File tree

1 file changed

+67
-13
lines changed

1 file changed

+67
-13
lines changed

src/cotlette/core/management/commands/createsuperuser.py

Lines changed: 67 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,15 @@ def handle(self, **options):
3434
email = options.get("email")
3535
noinput = options.get("noinput")
3636

37+
# Импортируем все модели для регистрации в ModelMeta
38+
self._import_all_models()
39+
40+
# Проверяем регистрацию моделей
41+
from cotlette.core.database.models import ModelMeta
42+
self.stdout.write(f'Зарегистрировано моделей: {len(ModelMeta._registry)}')
43+
for model_name, model in ModelMeta._registry.items():
44+
self.stdout.write(f' - {model_name}: {model.objects}')
45+
3746
# Try to get the User model
3847
try:
3948
from cotlette.contrib.auth.users.models import UserModel
@@ -47,6 +56,14 @@ def handle(self, **options):
4756
"Could not import User model. Make sure your apps are properly configured."
4857
)
4958

59+
# Try to get the Group model
60+
try:
61+
from cotlette.contrib.auth.groups.models import GroupModel
62+
if not GroupModel:
63+
GroupModel = None
64+
except ImportError:
65+
GroupModel = None
66+
5067
# Create the superuser
5168
if noinput:
5269
if not username:
@@ -66,9 +83,9 @@ def handle(self, **options):
6683
# Run async operations
6784
import asyncio
6885

69-
async def create_superuser_async():
86+
def create_superuser_sync():
7087
# Check if user already exists
71-
existing_user = await UserModel.objects.filter(email=email).first()
88+
existing_user = UserModel.objects.filter(email=email).first()
7289
if existing_user:
7390
raise CommandError(f"User with email '{email}' already exists.")
7491

@@ -82,15 +99,21 @@ async def create_superuser_async():
8299
hashed_password = hashlib.sha256(password.encode('utf-8')).hexdigest()
83100

84101
# Get or create group
85-
from cotlette.contrib.auth.groups.models import GroupModel
86102
if GroupModel:
87-
group = await GroupModel.objects.filter(id=1).first()
103+
group = GroupModel.objects.filter(id=1).first()
88104
if not group:
89-
group = await GroupModel.objects.create(
90-
name="Admin",
91-
description="Administrator group"
92-
)
93-
group_id = group.id
105+
# Создаем группу через ORM
106+
try:
107+
group = GroupModel.objects.create(
108+
name="Admin",
109+
description="Administrator group"
110+
)
111+
group_id = group.id if group else 1
112+
except Exception as e:
113+
# Группа уже существует или другая ошибка
114+
group_id = 1
115+
else:
116+
group_id = group.id
94117
else:
95118
group_id = 1
96119

@@ -119,17 +142,48 @@ async def create_superuser_async():
119142
# Выполняем запрос напрямую через базу данных
120143
from cotlette.core.database.sqlalchemy import db
121144
try:
122-
await db.execute_async(insert_query)
145+
db.execute(insert_query)
123146
except Exception as e:
124147
raise CommandError(f"Error creating superuser: {e}")
125148

126149
self.stdout.write(f"Superuser '{username}' created successfully.")
127150
except Exception as e:
128151
raise CommandError(f"Error creating superuser: {e}")
129152

130-
# Run the async function
131-
loop = asyncio.get_event_loop()
132-
loop.run_until_complete(create_superuser_async())
153+
# Run the sync function
154+
# Временно переопределяем should_use_async для использования синхронных операций
155+
import cotlette.core.database.query
156+
original_should_use_async = cotlette.core.database.query.should_use_async
157+
158+
def force_sync():
159+
return False
160+
161+
cotlette.core.database.query.should_use_async = force_sync
162+
163+
try:
164+
create_superuser_sync()
165+
finally:
166+
# Восстанавливаем оригинальную функцию
167+
cotlette.core.database.query.should_use_async = original_should_use_async
168+
169+
def _import_all_models(self):
170+
"""Импортирует все модели для регистрации в ModelMeta"""
171+
try:
172+
# Импортируем модели из contrib
173+
import cotlette.contrib.auth.users.models
174+
import cotlette.contrib.auth.groups.models
175+
import cotlette.contrib.admin.models
176+
177+
# Импортируем модели из приложений проекта
178+
try:
179+
import apps.home.models
180+
except ImportError:
181+
pass
182+
183+
except ImportError as e:
184+
self.stdout.write(
185+
self.style.WARNING(f'Предупреждение: не удалось импортировать некоторые модели: {e}')
186+
)
133187

134188
def _get_username(self, username=None):
135189
"""Get username from user input."""

0 commit comments

Comments
 (0)