Add a PRIMARY KEY lint

1 files changed, 134 insertions, 1 deletions

diff --git a/tests/storage/test_database.py b/tests/storage/test_database.py
index d60176b1d4..97e21fca15 100644
--- a/tests/storage/test_database.py
+++ b/tests/storage/test_database.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, List, Tuple
+from typing import Callable, List, Set, Tuple
 from unittest.mock import Mock, call
 
 from twisted.internet import defer
@@ -361,3 +361,136 @@ class PostgresReplicaIdentityTestCase(unittest.HomeserverTestCase):
                 0,
                 f"The following tables in the {pool.name()!r} database are missing REPLICA IDENTITIES: {missing!r}.",
             )
+
+
+class PostgresPrimaryKeyTestCase(unittest.HomeserverTestCase):
+    if not USE_POSTGRES_FOR_TESTS:
+        skip = "Requires Postgres"
+
+    # Some tables are exempt, mostly for historical reasons
+    EXEMPT_TABLES = {
+        "appservice_room_list",
+        "batch_events",
+        "blocked_rooms",
+        "cache_invalidation_stream_by_instance",
+        "current_state_delta_stream",
+        "deleted_pushers",
+        "device_auth_providers",
+        "device_federation_inbox",
+        "device_federation_outbox",
+        "device_inbox",
+        "device_lists_changes_in_room",
+        "device_lists_outbound_last_success",
+        "device_lists_outbound_pokes",
+        "device_lists_remote_cache",
+        "device_lists_remote_extremeties",
+        "device_lists_remote_resync",
+        "device_lists_stream",
+        "e2e_cross_signing_keys",
+        "e2e_cross_signing_signatures",
+        "e2e_room_keys",
+        "e2e_room_keys_versions",
+        "erased_users",
+        "event_auth",
+        "event_auth_chain_links",
+        "event_push_actions_staging",
+        "event_push_summary",
+        "event_relations",
+        "event_search",
+        "federation_inbound_events_staging",
+        "federation_stream_position",
+        "ignored_users",
+        "insertion_event_edges",
+        "insertion_event_extremities",
+        "insertion_events",
+        "local_media_repository_thumbnails",
+        "local_media_repository_url_cache",
+        "monthly_active_users",
+        "presence_stream",
+        "push_rules_stream",
+        "ratelimit_override",
+        "remote_media_cache_thumbnails",
+        "room_alias_servers",
+        "room_stats_earliest_token",
+        "room_stats_state",
+        "state_group_edges",
+        "state_groups_state",
+        "stream_ordering_to_exterm",
+        "stream_positions",
+        "threepid_guest_access_tokens",
+        "timeline_gaps",
+        "user_daily_visits",
+        "user_directory",
+        "user_directory_search",
+        "user_filters",
+        "user_ips",
+        "user_signature_stream",
+        "user_threepid_id_server",
+        "users_in_public_rooms",
+        "users_pending_deactivation",
+        "users_who_share_private_rooms",
+        "worker_locks",
+    }
+
+    def prepare(
+        self, reactor: MemoryReactor, clock: Clock, homeserver: HomeServer
+    ) -> None:
+        self.db_pools = homeserver.get_datastores().databases
+
+    def test_all_nonexempt_tables_have_primary_key(self) -> None:
+        """
+        Tests that all tables have a PRIMARY KEY.
+
+        Whilst there are occasionally potentially good reasons to not have one,
+        we should default to having PRIMARY KEYs on new tables.
+
+        Historically we have not created PRIMARY KEYs when they should be created
+        (sometimes even creating a UNIQUE index on the same columns instead).
+
+        It is preferable to use PRIMARY KEYs instead of UNIQUE indices where possible
+        because this:
+
+        - provides a default REPLICA IDENTITY, ensuring logical replication works
+          without extra complication;
+        - is sometimes more helpful for understanding; and
+        - is more useful to external tools which expect PRIMARY KEYs (e.g. some database
+          migration tools on some cloud providers do not work with tables without PRIMARY
+          KEYs)
+        """
+
+        sql = """
+            SELECT tbl.table_name
+            FROM information_schema.tables tbl
+            WHERE table_type = 'BASE TABLE'
+                AND table_schema not in ('pg_catalog', 'information_schema')
+                AND NOT EXISTS (
+                    SELECT 1
+                    FROM information_schema.key_column_usage kcu
+                    WHERE kcu.table_name = tbl.table_name
+                        AND kcu.table_schema = tbl.table_schema
+                )
+        """
+
+        def _list_tables_with_missing_primary_keys_txn(
+            txn: LoggingTransaction,
+        ) -> Set[str]:
+            txn.execute(sql)
+            return {table_name for table_name, in txn}
+
+        for pool in self.db_pools:
+            missing = sorted(
+                self.get_success(
+                    pool.runInteraction(
+                        "test_list_missing_primary_keys",
+                        _list_tables_with_missing_primary_keys_txn,
+                    )
+                )
+                - self.EXEMPT_TABLES
+            )
+            self.assertEqual(
+                len(missing),
+                0,
+                f"The following tables in the {pool.name()!r} database are missing PRIMARY KEYs: {missing!r}.\n"
+                f"Please consider adding a primary key if it makes sense to do so — this gives the table a replica identity for free and helps humans and tools to understand the table better.\n"
+                f"If this lack of a primary key is intentional, please add an exemption.",
+            )