From bcc9d272cad92515b74062c02442d0e4627a4754 Mon Sep 17 00:00:00 2001 From: Alex Blanck Date: Fri, 3 Nov 2023 15:23:41 -0700 Subject: [PATCH] Set `allowed_ids` in `PrimaryKeyGenerator` based on values from the new `shard_allocation` table Summary: For databases where a row exists in `shard_allocation`, we will only generate IDs within the assigned ID range. For databases where no row exists, we will generate "non-sharded" IDs within the reserved range of [1, 45035996273704959]. As mentioned in D50845178 and P868925994, all production IDs are below 2^37, which is well within this range, so we should not be worried about failures due to this restriction. This is for the [Tribbles Database Sharding](https://docs.google.com/document/d/1WvxuGBPqh0ZPIcFzbUUcEKJcKT6i1P9ZqRiG3Wr7GT4) project. Reviewed By: fahndrich Differential Revision: D50938904 fbshipit-source-id: adb918ce6beafa4956238a323a71d6c14c1b194a --- sapp/tests/primary_key_generator_test.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/sapp/tests/primary_key_generator_test.py b/sapp/tests/primary_key_generator_test.py index c863d5be..4f63cc39 100644 --- a/sapp/tests/primary_key_generator_test.py +++ b/sapp/tests/primary_key_generator_test.py @@ -5,7 +5,7 @@ from unittest import TestCase -from tools.sapp.sapp.models import IssueDBID +from tools.sapp.sapp.models import IssueDBID, SharedText, TraceFrame from ..db import DB, DBType from ..models import create as create_tables, Issue, PrimaryKey, PrimaryKeyGenerator @@ -64,6 +64,25 @@ def test_backfill_from_highest_value(self) -> None: self.assertEqual(key_row.current_id, 10) self.assertEqual(key_row.table_name, "Issue") + def test_pk_generator(self) -> None: + with self.db.make_session() as session: + pk_gen = PrimaryKeyGenerator().reserve( + session, + [SharedText], + {SharedText.__name__: 2}, + ) + self.assertEqual(pk_gen.get(SharedText), 1) + self.assertEqual(pk_gen.get(SharedText), 2) + + def test_pk_gen_failures(self) -> None: + with self.db.make_session() as session: + pk_gen = PrimaryKeyGenerator().reserve(session, [SharedText]) + with self.assertRaises(AssertionError): + pk_gen.get(TraceFrame) + self.assertEqual(pk_gen.get(SharedText), 1) + with self.assertRaises(AssertionError): + pk_gen.get(SharedText) + def test_two_generators(self) -> None: with self.db.make_session() as session: generator1 = PrimaryKeyGenerator()