Set allowed_ids in PrimaryKeyGenerator based on values from the n…

…ew `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

@@ -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()