Repository: zulip/zulip

Issue:
Use single database query in `do_incremental_logging_stat`
We use an inefficient database query construction for `do_incremental_logging_stat`:

```
    row, created = table._default_manager.get_or_create(
        property=stat.property,
        subgroup=subgroup,
        end_time=end_time,
        defaults={"value": increment},
        **id_args,
    )
    if not created:
        row.value = F("value") + increment
        row.save(update_fields=["value"])
```

Django (and postgres) support doing this operation in a single SQL query, perhaps using https://docs.djangoproject.com/en/5.0/ref/models/expressions/#f-expressions. I'm not sure whether https://docs.djangoproject.com/en/5.0/ref/models/querysets/#update-or-create will do the right thing in Django 5, or if it'll do multiple queries.

You can use this settings tweak:
```
diff --git a/zproject/computed_settings.py b/zproject/computed_settings.py
index 8d691753da..813a896956 100644
--- a/zproject/computed_settings.py
+++ b/zproject/computed_settings.py
@@ -864,11 +864,11 @@ LOGGING: Dict[str, Any] = {
             "propagate": False,
         },
         ## Uncomment the following to get all database queries logged to the console
-        # 'django.db': {
-        #     'level': 'DEBUG',
-        #     'handlers': ['console'],
-        #     'propagate': False,
-        # },
+        "django.db": {
+            "level": "DEBUG",
+            "handlers": ["console"],
+            "propagate": False,
+        },
         # other libraries, alphabetized
         "django_auth_ldap": {
             "level": "DEBUG",
```

to make Django print all database queries to verify your work.

See [this chat.zulip.org thread](https://chat.zulip.org/#narrow/stream/3-backend/topic/mark-as-unread.20performance/near/1738515) for details.


Discussion:
Hello @zulip/server-analytics members, this issue was labeled with the "area: analytics" label, so you may want to check it out!

<!-- areaLabelAddition -->

@zulipbot claim
@Mahhheshh You have been unassigned from this issue because you have not made any updates for over 14 days. Please feel free to reclaim the issue if you decide to pick up again. Thanks!

yes
Can I work on this issue?
@zulipbot claim
Hello @ritikraj26, it looks like you've currently claimed 1 issue in this repository. We encourage new contributors to focus their efforts on at most 1 issue at a time, so please complete your work on your other claimed issues before trying to claim this issue again.

We look forward to your valuable contributions!


Resolving patch:
diff --git a/analytics/lib/counts.py b/analytics/lib/counts.py
--- a/analytics/lib/counts.py
+++ b/analytics/lib/counts.py
@@ -2,11 +2,10 @@
 import time
 from collections import OrderedDict, defaultdict
 from datetime import datetime, timedelta
-from typing import Callable, Dict, Optional, Sequence, Tuple, Type, Union
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 from django.conf import settings
 from django.db import connection, models
-from django.db.models import F
 from psycopg2.sql import SQL, Composable, Identifier, Literal
 from typing_extensions import TypeAlias, override
 
@@ -312,27 +311,47 @@ def do_increment_logging_stat(
         return
 
     table = stat.data_collector.output_table
+    id_args: Dict[str, Union[int, None]] = {}
+    conflict_args: List[str] = []
     if table == RealmCount:
         assert isinstance(model_object_for_bucket, Realm)
-        id_args: Dict[
-            str, Optional[Union[Realm, UserProfile, Stream, "RemoteRealm", "RemoteZulipServer"]]
-        ] = {"realm": model_object_for_bucket}
+        id_args = {"realm_id": model_object_for_bucket.id}
+        conflict_args = ["realm_id"]
     elif table == UserCount:
         assert isinstance(model_object_for_bucket, UserProfile)
-        id_args = {"realm": model_object_for_bucket.realm, "user": model_object_for_bucket}
+        id_args = {
+            "realm_id": model_object_for_bucket.realm_id,
+            "user_id": model_object_for_bucket.id,
+        }
+        conflict_args = ["user_id"]
     elif table == StreamCount:
         assert isinstance(model_object_for_bucket, Stream)
-        id_args = {"realm": model_object_for_bucket.realm, "stream": model_object_for_bucket}
+        id_args = {
+            "realm_id": model_object_for_bucket.realm_id,
+            "stream_id": model_object_for_bucket.id,
+        }
+        conflict_args = ["stream_id"]
     elif table == RemoteInstallationCount:
         assert isinstance(model_object_for_bucket, RemoteZulipServer)
-        id_args = {"server": model_object_for_bucket, "remote_id": None}
+        id_args = {"server_id": model_object_for_bucket.id, "remote_id": None}
+        conflict_args = ["server_id"]
     elif table == RemoteRealmCount:
         assert isinstance(model_object_for_bucket, RemoteRealm)
+        # For RemoteRealmCount (e.g. `mobile_pushes_forwarded::day`),
+        # we have no `remote_id` nor `realm_id`, since they are not
+        # imported from the remote server, which is the source of
+        # truth of those two columns.  Their "ON CONFLICT" is thus the
+        # only unique key we have, which is `remote_realm_id`, and not
+        # `server_id` / `realm_id`.
         id_args = {
-            "server": model_object_for_bucket.server,
-            "remote_realm": model_object_for_bucket,
+            "server_id": model_object_for_bucket.server_id,
+            "remote_realm_id": model_object_for_bucket.id,
             "remote_id": None,
+            "realm_id": None,
         }
+        conflict_args = [
+            "remote_realm_id",
+        ]
     else:
         raise AssertionError("Unsupported CountStat output_table")
 
@@ -343,16 +362,49 @@ def do_increment_logging_stat(
     else:
         raise AssertionError("Unsupported CountStat frequency")
 
-    row, created = table._default_manager.get_or_create(
-        property=stat.property,
-        subgroup=subgroup,
-        end_time=end_time,
-        defaults={"value": increment},
-        **id_args,
+    is_subgroup: SQL = SQL("NULL")
+    if subgroup is not None:
+        is_subgroup = SQL("NOT NULL")
+        # For backwards consistency, we cast the subgroup to a string
+        # in Python; this emulates the behaviour of `get_or_create`,
+        # which was previously used in this function, and performed
+        # this cast because the `subgroup` column is defined as a
+        # `CharField`.  Omitting this explicit cast causes a subgroup
+        # of the boolean False to be passed as the PostgreSQL false,
+        # which it stringifies as the lower-case `'false'`, not the
+        # initial-case `'False'` if Python stringifies it.
+        #
+        # Other parts of the system (e.g. count_message_by_user_query)
+        # already use PostgreSQL to cast bools to strings, resulting
+        # in `subgroup` values of lower-case `'false'` -- for example
+        # in `messages_sent:is_bot:hour`.  Fixing this inconsistency
+        # via a migration is complicated by these records being
+        # exchanged over the wire from remote servers.
+        subgroup = str(subgroup)
+        conflict_args.append("subgroup")
+
+    id_column_names = SQL(", ").join(map(Identifier, id_args.keys()))
+    id_values = SQL(", ").join(map(Literal, id_args.values()))
+    conflict_columns = SQL(", ").join(map(Identifier, conflict_args))
+
+    sql_query = SQL(
+        """
+        INSERT INTO {table_name}(property, subgroup, end_time, value, {id_column_names})
+        VALUES (%s, %s, %s, %s, {id_values})
+        ON CONFLICT (property, end_time, {conflict_columns})
+        WHERE subgroup IS {is_subgroup}
+        DO UPDATE SET
+            value = {table_name}.value + EXCLUDED.value
+        """
+    ).format(
+        table_name=Identifier(table._meta.db_table),
+        id_column_names=id_column_names,
+        id_values=id_values,
+        conflict_columns=conflict_columns,
+        is_subgroup=is_subgroup,
     )
-    if not created:
-        row.value = F("value") + increment
-        row.save(update_fields=["value"])
+    with connection.cursor() as cursor:
+        cursor.execute(sql_query, [stat.property, subgroup, end_time, increment])
 
 
 def do_drop_all_analytics_tables() -> None:
diff --git a/zilencer/migrations/0060_remove_remoterealmcount_unique_remote_realm_installation_count_and_more.py b/zilencer/migrations/0060_remove_remoterealmcount_unique_remote_realm_installation_count_and_more.py
new file mode 100644
--- /dev/null
+++ b/zilencer/migrations/0060_remove_remoterealmcount_unique_remote_realm_installation_count_and_more.py
@@ -0,0 +1,81 @@
+# Generated by Django 5.0.5 on 2024-05-06 15:17
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("zilencer", "0059_remoterealmauditlog_add_synced_billing_event_type_index"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            # Django does not provide a RenameConstraint operation.
+            # "Constraints" are created as indexes in PostgreSQL, so
+            # we rename the underlying indexes.
+            database_operations=[
+                migrations.RunSQL(
+                    sql=(
+                        "ALTER INDEX unique_remote_realm_installation_count "
+                        "RENAME TO unique_server_realm_installation_count"
+                    ),
+                    reverse_sql=(
+                        "ALTER INDEX unique_server_realm_installation_count "
+                        "RENAME TO unique_remote_realm_installation_count"
+                    ),
+                ),
+                migrations.RunSQL(
+                    sql=(
+                        "ALTER INDEX unique_remote_realm_installation_count_null_subgroup "
+                        "RENAME TO unique_server_realm_installation_count_null_subgroup"
+                    ),
+                    reverse_sql=(
+                        "ALTER INDEX unique_server_realm_installation_count_null_subgroup "
+                        "RENAME TO unique_remote_realm_installation_count_null_subgroup"
+                    ),
+                ),
+            ],
+            state_operations=[
+                migrations.RemoveConstraint(
+                    model_name="remoterealmcount",
+                    name="unique_remote_realm_installation_count",
+                ),
+                migrations.RemoveConstraint(
+                    model_name="remoterealmcount",
+                    name="unique_remote_realm_installation_count_null_subgroup",
+                ),
+                migrations.AddConstraint(
+                    model_name="remoterealmcount",
+                    constraint=models.UniqueConstraint(
+                        condition=models.Q(("subgroup__isnull", False)),
+                        fields=("server", "realm_id", "property", "subgroup", "end_time"),
+                        name="unique_server_realm_installation_count",
+                    ),
+                ),
+                migrations.AddConstraint(
+                    model_name="remoterealmcount",
+                    constraint=models.UniqueConstraint(
+                        condition=models.Q(("subgroup__isnull", True)),
+                        fields=("server", "realm_id", "property", "end_time"),
+                        name="unique_server_realm_installation_count_null_subgroup",
+                    ),
+                ),
+            ],
+        ),
+        migrations.AddConstraint(
+            model_name="remoterealmcount",
+            constraint=models.UniqueConstraint(
+                condition=models.Q(("subgroup__isnull", False)),
+                fields=("remote_realm_id", "property", "subgroup", "end_time"),
+                name="unique_remote_realm_installation_count",
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="remoterealmcount",
+            constraint=models.UniqueConstraint(
+                condition=models.Q(("subgroup__isnull", True)),
+                fields=("remote_realm_id", "property", "end_time"),
+                name="unique_remote_realm_installation_count_null_subgroup",
+            ),
+        ),
+    ]
diff --git a/zilencer/models.py b/zilencer/models.py
--- a/zilencer/models.py
+++ b/zilencer/models.py
@@ -416,14 +416,28 @@ class RemoteRealmCount(BaseRemoteCount):
 
     class Meta:
         constraints = [
+            # These two constraints come from the information as
+            # provided by the remote server, for rows they provide.
             UniqueConstraint(
                 fields=["server", "realm_id", "property", "subgroup", "end_time"],
                 condition=Q(subgroup__isnull=False),
-                name="unique_remote_realm_installation_count",
+                name="unique_server_realm_installation_count",
             ),
             UniqueConstraint(
                 fields=["server", "realm_id", "property", "end_time"],
                 condition=Q(subgroup__isnull=True),
+                name="unique_server_realm_installation_count_null_subgroup",
+            ),
+            # These two constraints come from our internal
+            # record-keeping, which has a RemoteRealm object.
+            UniqueConstraint(
+                fields=["remote_realm_id", "property", "subgroup", "end_time"],
+                condition=Q(subgroup__isnull=False),
+                name="unique_remote_realm_installation_count",
+            ),
+            UniqueConstraint(
+                fields=["remote_realm_id", "property", "end_time"],
+                condition=Q(subgroup__isnull=True),
                 name="unique_remote_realm_installation_count_null_subgroup",
             ),
             UniqueConstraint(


Test patch:
diff --git a/analytics/tests/test_counts.py b/analytics/tests/test_counts.py
--- a/analytics/tests/test_counts.py
+++ b/analytics/tests/test_counts.py
@@ -1324,6 +1324,11 @@ def test_increment(self) -> None:
         do_increment_logging_stat(self.default_realm, stat, None, self.TIME_ZERO)
         self.assertTableState(RealmCount, ["value"], [[3]])
 
+    def test_do_increment_logging_start_query_count(self) -> None:
+        stat = LoggingCountStat("test", RealmCount, CountStat.DAY)
+        with self.assert_database_query_count(1):
+            do_increment_logging_stat(self.default_realm, stat, None, self.TIME_ZERO)
+
 
 class TestLoggingCountStats(AnalyticsTestCase):
     def test_aggregation(self) -> None:
diff --git a/zerver/tests/test_signup.py b/zerver/tests/test_signup.py
--- a/zerver/tests/test_signup.py
+++ b/zerver/tests/test_signup.py
@@ -932,7 +932,7 @@ def test_register(self) -> None:
         # seem to be any O(N) behavior.  Some of the cache hits are related
         # to sending messages, such as getting the welcome bot, looking up
         # the alert words for a realm, etc.
-        with self.assert_database_query_count(105), self.assert_memcached_count(18):
+        with self.assert_database_query_count(104), self.assert_memcached_count(18):
             with self.captureOnCommitCallbacks(execute=True):
                 self.register(self.nonreg_email("test"), "test")
 
diff --git a/zerver/tests/test_users.py b/zerver/tests/test_users.py
--- a/zerver/tests/test_users.py
+++ b/zerver/tests/test_users.py
@@ -907,7 +907,7 @@ def test_create_user_with_multiple_streams(self) -> None:
 
         prereg_user = PreregistrationUser.objects.get(email="fred@zulip.com")
 
-        with self.assert_database_query_count(94):
+        with self.assert_database_query_count(93):
             with self.assert_memcached_count(23):
                 with self.capture_send_event_calls(expected_num_events=11) as events:
                     fred = do_create_user(