Merge pull request #12 from ali-zahedi/feat/multi-accounts

feat: support multi accounts

Author: ali-zahedi

django_telethon/admin.py

@@ -32,6 +32,7 @@ class LoginAdmin(admin.ModelAdmin):
         'created_at',
     ]
     list_filter = ['created_at', 'client_session__name']
+    raw_id_fields = ('client_session',)
 
 
 @admin.register(Session)
@@ -92,8 +93,9 @@ def send_a_test_message(self, request, queryset):
 @admin.register(UpdateState)
 class UpdateStateAdmin(admin.ModelAdmin):
     list_display = [
-        'id',
+        'pk',
         'client_session',
+        'entity_id',
         'pts',
         'qts',
         'date',
@@ -102,6 +104,14 @@ class UpdateStateAdmin(admin.ModelAdmin):
     raw_id_fields = ['client_session']
     list_filter = ['client_session__name']
 
+    @admin.display(description="client session")
+    def client_session(self, obj):
+        return obj.client_session.name
+
+    def get_queryset(self, request):
+        queryset = super().get_queryset(request)
+        return queryset.select_related('client_session')
+
 
 @admin.register(SentFile)
 class SentFileAdmin(admin.ModelAdmin):

django_telethon/management/commands/runtelegram.py

@@ -16,13 +16,13 @@ async def _run_forever():
         logging.exception(e, exc_info=True)
 
     while True:
-        await asyncio.sleep(30)
         try:
             await re_connect_clients()
         except KeyboardInterrupt:
             break
         except Exception as e:
             logging.exception(e, exc_info=True)
+        await asyncio.sleep(30)
 
 
 async def _main():

django_telethon/migrations/0002_updatestate_entity.py

@@ -0,0 +1,33 @@
+# Generated by Django 5.0.3 on 2024-12-03 14:01
+
+from django.db import migrations, models
+
+
+def _migrate_from_old_column(apps, schema_editor) -> None:
+    UpdateState = apps.get_model("django_telethon", "UpdateState")
+    i = 0
+    for update_state in UpdateState.objects.all():
+        i += 1
+        UpdateState.objects.filter(id=update_state.id).update(id=i, entity_id=update_state.id)
+
+    # Step 2: Reset primary keys
+    with schema_editor.connection.cursor() as cursor:
+        cursor.execute(f"ALTER SEQUENCE {UpdateState._meta.db_table}_id_seq RESTART WITH {i +  1}")
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('django_telethon', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='updatestate',
+            name='entity_id',
+            field=models.BigIntegerField(null=True),
+        ),
+        migrations.RunPython(
+            _migrate_from_old_column,
+            reverse_code=lambda x, y: None,
+        ),
+    ]

django_telethon/migrations/0003_alter_updatestate_entity_id.py

@@ -0,0 +1,21 @@
+# Generated by Django 5.0.3 on 2024-12-03 17:53
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('django_telethon', '0002_updatestate_entity'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='updatestate',
+            name='entity_id',
+            field=models.BigIntegerField(db_index=True),
+        ),
+        migrations.AlterUniqueTogether(
+            name='updatestate',
+            unique_together={('client_session', 'entity_id')},
+        ),
+    ]

django_telethon/models/states.py

@@ -12,6 +12,8 @@ class UpdateState(models.Model):
         on_delete=models.CASCADE,
         verbose_name=_('Client Session'),
     )
+    entity_id = models.BigIntegerField(db_index=True)
+
     pts = models.IntegerField(
         verbose_name=_('pts'),
     )
@@ -26,5 +28,6 @@ class UpdateState(models.Model):
     )
 
     class Meta:
+        unique_together = (('client_session', 'entity_id'),)
         verbose_name = _('Update state')
         verbose_name_plural = _('Update states')

django_telethon/sessions.py

@@ -97,14 +97,14 @@ def _update_session_table(self):
 
     def get_update_state(self, entity_id):
         try:
-            state = self.client_session.updatestate_set.get(pk=entity_id)
+            state = self.client_session.updatestate_set.get(entity_id=entity_id)
             return types.updates.State(state.pts, state.qts, state.date, state.seq, unread_count=0)
         except UpdateState.DoesNotExist:
             return None
 
     def set_update_state(self, entity_id, state):
         self.client_session.updatestate_set.update_or_create(
-            pk=entity_id,
+            entity_id=entity_id,
             defaults={
                 'pts': state.pts,
                 'qts': state.qts,

django_telethon/utils.py

@@ -51,7 +51,9 @@ async def connect_client(client_app, app):
     await telegram_client.connect()
     if not await telegram_client.is_user_authorized():
         client_app.login_status = LoginStatus.LOGIN_REQUIRED
-        client_app.save()
+        client_app.save(update_fields=['login_status'])
+        if client_app.session:
+            client_app.session.delete()
         logging.critical(f"Authorization failed for client: {client_app.name}")
         return
     if client_app.login_status != LoginStatus.LOGIN_DONE: