From bcc69494619953dcc528a8fc8f5501de9a0f31d6 Mon Sep 17 00:00:00 2001
From: Tim Abbott <tabbott@zulipchat.com>
Date: Tue, 23 Apr 2019 13:32:12 -0700
Subject: [PATCH] zilencer: Add better error handling for IntegrityError.

This provides a clean warning and 40x error, rather than a 500, for
this corner case which is very likely user error.

The test here is awkward because we have to work around
https://github.com/zulip/zulip/issues/12362.
---
 zerver/tests/test_push_notifications.py | 19 ++++++++++++++++++-
 zilencer/views.py                       | 15 +++++++++++++--
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/zerver/tests/test_push_notifications.py b/zerver/tests/test_push_notifications.py
index a43c2fd1fa..161607e564 100644
--- a/zerver/tests/test_push_notifications.py
+++ b/zerver/tests/test_push_notifications.py
@@ -15,6 +15,7 @@ import uuid
 from django.test import override_settings
 from django.conf import settings
 from django.http import HttpResponse
+from django.db import transaction
 from django.db.models import F
 from django.utils.crypto import get_random_string
 from django.utils.timezone import utc as timezone_utc
@@ -322,7 +323,10 @@ class AnalyticsBouncerTest(BouncerTestCase):
         RealmCount.objects.create(
             realm=user.realm, property=realm_stat.property, end_time=end_time, value=5)
         InstallationCount.objects.create(
-            property=realm_stat.property, end_time=end_time, value=5)
+            property=realm_stat.property, end_time=end_time, value=5,
+            # We set a subgroup here to work around:
+            # https://github.com/zulip/zulip/issues/12362
+            subgroup="test_subgroup")
 
         self.assertEqual(RealmCount.objects.count(), 1)
         self.assertEqual(InstallationCount.objects.count(), 1)
@@ -368,6 +372,19 @@ class AnalyticsBouncerTest(BouncerTestCase):
                                subdomain="")
         self.assert_json_error(result, "Data is out of order.")
 
+        with mock.patch("zilencer.views.validate_count_stats"):
+            # We need to wrap a transaction here to avoid the
+            # IntegrityError that will be thrown in here from breaking
+            # the unittest transaction.
+            with transaction.atomic():
+                result = self.api_post(
+                    self.server_uuid,
+                    '/api/v1/remotes/server/analytics',
+                    {'realm_counts': ujson.dumps(realm_count_data),
+                     'installation_counts': ujson.dumps(installation_count_data)},
+                    subdomain="")
+            self.assert_json_error(result, "Invalid data.")
+
     @override_settings(PUSH_NOTIFICATION_BOUNCER_URL='https://push.zulip.org.example.com')
     @mock.patch('zerver.lib.push_notifications.requests.request')
     def test_analytics_api_invalid(self, mock_request: Any) -> None:
diff --git a/zilencer/views.py b/zilencer/views.py
index 5045906bbf..6000508d05 100644
--- a/zilencer/views.py
+++ b/zilencer/views.py
@@ -1,5 +1,6 @@
 from typing import Any, Dict, List, Optional, Union, cast
 import datetime
+import logging
 
 from django.core.exceptions import ValidationError
 from django.core.validators import validate_email, URLValidator
@@ -201,7 +202,12 @@ def remote_server_post_analytics(request: HttpRequest,
                 end_time=datetime.datetime.fromtimestamp(item['end_time'], tz=timezone_utc),
                 subgroup=item['subgroup'],
                 value=item['value']))
-        RemoteRealmCount.objects.bulk_create(objects_to_create)
+        try:
+            RemoteRealmCount.objects.bulk_create(objects_to_create)
+        except IntegrityError:
+            logging.warning("Invalid data saving RemoteRealmCount for server %s/%s" % (
+                server.hostname, server.uuid))
+            return json_error(_("Invalid data."))
 
     while len(installation_counts) > 0:
         batch = installation_counts[0:BATCH_SIZE]
@@ -216,7 +222,12 @@ def remote_server_post_analytics(request: HttpRequest,
                 end_time=datetime.datetime.fromtimestamp(item['end_time'], tz=timezone_utc),
                 subgroup=item['subgroup'],
                 value=item['value']))
-        RemoteInstallationCount.objects.bulk_create(objects_to_create)
+        try:
+            RemoteInstallationCount.objects.bulk_create(objects_to_create)
+        except IntegrityError:
+            logging.warning("Invalid data saving RemoteInstallationCount for server %s/%s" % (
+                server.hostname, server.uuid))
+            return json_error(_("Invalid data."))
     return json_success()
 
 def get_last_id_from_server(server: RemoteZulipServer, model: Any) -> int: