diff --git a/zerver/lib/export.py b/zerver/lib/export.py index 755d1c4c6b..e545676789 100644 --- a/zerver/lib/export.py +++ b/zerver/lib/export.py @@ -122,6 +122,7 @@ ALL_ZULIP_TABLES = { 'zerver_userprofile', 'zerver_userprofile_groups', 'zerver_userprofile_user_permissions', + 'zerver_userstatus', 'zerver_mutedtopic', } @@ -191,6 +192,9 @@ NON_EXPORTED_TABLES = { 'zerver_defaultstreamgroup_streams', 'zerver_submessage', + # This is low priority, since users can easily just reset themselves to away. + 'zerver_userstatus', + # For any tables listed below here, it's a bug that they are not present in the export. } diff --git a/zerver/lib/user_status.py b/zerver/lib/user_status.py new file mode 100644 index 0000000000..9a1ecc6c28 --- /dev/null +++ b/zerver/lib/user_status.py @@ -0,0 +1,36 @@ +from django.utils.timezone import now as timezone_now + +from zerver.models import ( + UserStatus, +) + +from typing import Set + +def get_away_user_ids(realm_id: int) -> Set[int]: + user_ids = UserStatus.objects.filter( + status=UserStatus.AWAY, + user_profile__realm_id=realm_id, + user_profile__is_active=True, + ).values_list('user_profile_id', flat=True) + + return set(user_ids) + +def set_away_status(user_profile_id: int, + client_id: int) -> None: + + timestamp = timezone_now() + status = UserStatus.AWAY + + UserStatus.objects.update_or_create( + user_profile_id=user_profile_id, + defaults=dict( + client_id=client_id, + timestamp=timestamp, + status=status, + ), + ) + +def revoke_away_status(user_profile_id: int) -> None: + UserStatus.objects.filter( + user_profile_id=user_profile_id, + ).delete() diff --git a/zerver/migrations/0199_userstatus.py b/zerver/migrations/0199_userstatus.py new file mode 100644 index 0000000000..6a8072037f --- /dev/null +++ b/zerver/migrations/0199_userstatus.py @@ -0,0 +1,27 @@ +# -*- coding: utf-8 -*- +# Generated by Django 1.11.16 on 2018-12-17 18:49 +from __future__ import unicode_literals + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('zerver', '0198_preregistrationuser_invited_as'), + ] + + operations = [ + migrations.CreateModel( + name='UserStatus', + fields=[ + ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('timestamp', models.DateTimeField()), + ('status', models.PositiveSmallIntegerField(default=1)), + ('client', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='zerver.Client')), + ('user_profile', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL)), + ], + ), + ] diff --git a/zerver/models.py b/zerver/models.py index d7f7d4f44b..48c2f2a77e 100644 --- a/zerver/models.py +++ b/zerver/models.py @@ -2170,6 +2170,16 @@ class UserPresence(models.Model): return status_val +class UserStatus(models.Model): + user_profile = models.OneToOneField(UserProfile, on_delete=CASCADE) # type: UserProfile + + timestamp = models.DateTimeField() # type: datetime.datetime + client = models.ForeignKey(Client, on_delete=CASCADE) # type: Client + + AWAY = 1 + + status = models.PositiveSmallIntegerField(default=AWAY) # type: int + class DefaultStream(models.Model): realm = models.ForeignKey(Realm, on_delete=CASCADE) # type: Realm stream = models.ForeignKey(Stream, on_delete=CASCADE) # type: Stream diff --git a/zerver/tests/test_user_status.py b/zerver/tests/test_user_status.py new file mode 100644 index 0000000000..5b84fbe2db --- /dev/null +++ b/zerver/tests/test_user_status.py @@ -0,0 +1,88 @@ +from zerver.lib.test_classes import ( + ZulipTestCase, +) +from zerver.lib.user_status import ( + get_away_user_ids, + revoke_away_status, + set_away_status, +) + +from zerver.models import ( + get_client, + UserStatus, +) + +class UserStatusTest(ZulipTestCase): + def test_basics(self) -> None: + cordelia = self.example_user('cordelia') + hamlet = self.example_user('hamlet') + king_lear = self.lear_user('king') + + realm_id = hamlet.realm_id + + away_user_ids = get_away_user_ids(realm_id=realm_id) + self.assertEqual(away_user_ids, set()) + + client1 = get_client('web') + client2 = get_client('ZT') + + set_away_status( + user_profile_id=hamlet.id, + client_id=client1.id, + ) + + away_user_ids = get_away_user_ids(realm_id=realm_id) + self.assertEqual(away_user_ids, {hamlet.id}) + + # Test that second client just updates + # the record. We only store one record + # per user. The user's status transcends + # clients; we only store the client for + # reference and to maybe reconcile timeout + # situations. + set_away_status( + user_profile_id=hamlet.id, + client_id=client2.id, + ) + + away_user_ids = get_away_user_ids(realm_id=realm_id) + self.assertEqual(away_user_ids, {hamlet.id}) + + rec_count = UserStatus.objects.filter(user_profile_id=hamlet.id).count() + self.assertEqual(rec_count, 1) + + revoke_away_status( + user_profile_id=hamlet.id, + ) + + away_user_ids = get_away_user_ids(realm_id=realm_id) + self.assertEqual(away_user_ids, set()) + + # Now set away status for three different users across + # two realms. + set_away_status( + user_profile_id=hamlet.id, + client_id=client1.id, + ) + set_away_status( + user_profile_id=cordelia.id, + client_id=client2.id, + ) + set_away_status( + user_profile_id=king_lear.id, + client_id=client2.id, + ) + + away_user_ids = get_away_user_ids(realm_id=realm_id) + self.assertEqual(away_user_ids, {cordelia.id, hamlet.id}) + + away_user_ids = get_away_user_ids(realm_id=king_lear.realm.id) + self.assertEqual(away_user_ids, {king_lear.id}) + + # Revoke Hamlet again. + revoke_away_status( + user_profile_id=hamlet.id, + ) + + away_user_ids = get_away_user_ids(realm_id=realm_id) + self.assertEqual(away_user_ids, {cordelia.id})