ref: Type all other bitfields in models (#51729)

migrations_lockfile.txt

@@ -6,5 +6,5 @@ To resolve this, rebase against latest master and regenerate your migration. Thi
 will then be regenerated, and you should be able to merge without conflicts.
 
 nodestore: 0002_nodestore_no_dictfield
-sentry: 0500_set_none_date_last_modified_to_date_uploaded
+sentry: 0501_typed_bitfield_remove_labels
 social_auth: 0001_initial

pyproject.toml

@@ -751,7 +751,6 @@ module = [
     "sentry.models.release",
     "sentry.models.releasefile",
     "sentry.models.team",
-    "sentry.models.user",
     "sentry.monitors.consumers.monitor_consumer",
     "sentry.monitors.endpoints.base",
     "sentry.monitors.endpoints.monitor_ingest_checkin_attachment",
@@ -972,7 +971,6 @@ module = [
     "sentry.tasks.servicehooks",
     "sentry.tasks.store",
     "sentry.tasks.unmerge",
-    "sentry.tasks.weekly_reports",
     "sentry.templatetags.sentry_admin_helpers",
     "sentry.templatetags.sentry_assets",
     "sentry.templatetags.sentry_helpers",
@@ -1150,7 +1148,6 @@ module = [
     "tests.sentry.api.endpoints.test_organization_details",
     "tests.sentry.api.endpoints.test_organization_index",
     "tests.sentry.api.endpoints.test_organization_member_details",
-    "tests.sentry.api.endpoints.test_organization_member_index",
     "tests.sentry.api.endpoints.test_organization_metric_data",
     "tests.sentry.api.endpoints.test_organization_metric_details",
     "tests.sentry.api.endpoints.test_organization_metric_tag_details",
@@ -1490,8 +1487,6 @@ module = [
     "tests.sentry.tasks.deletion.test_scheduled",
     "tests.sentry.tasks.integrations.github.test_pr_comment",
     "tests.sentry.tasks.test_assemble",
-    "tests.sentry.tasks.test_auth",
-    "tests.sentry.tasks.test_check_auth",
     "tests.sentry.tasks.test_derive_code_mappings",
     "tests.sentry.tasks.test_groupowner",
     "tests.sentry.tasks.test_merge",
@@ -1526,7 +1521,6 @@ module = [
     "tests.sentry.web.frontend.test_auth_oauth2",
     "tests.sentry.web.frontend.test_auth_organization_login",
     "tests.sentry.web.frontend.test_auth_saml2",
-    "tests.sentry.web.frontend.test_disabled_member_view",
     "tests.sentry.web.frontend.test_oauth_authorize",
     "tests.sentry.web.frontend.test_organization_auth_settings",
     "tests.sentry.web.frontend.test_react_page",

src/bitfield/__init__.py

@@ -1,6 +1,13 @@
-from bitfield.models import BitField, TypedBitfield  # NOQA
+from bitfield.models import BitField, TypedClassBitField, typed_dict_bitfield  # NOQA
 from bitfield.types import Bit, BitHandler
 
 default_app_config = "bitfield.apps.BitFieldAppConfig"
 
-__all__ = ("Bit", "BitField", "BitHandler", "default_app_config", "TypedBitfield")
+__all__ = (
+    "Bit",
+    "BitField",
+    "BitHandler",
+    "default_app_config",
+    "TypedClassBitField",
+    "typed_dict_bitfield",
+)

src/bitfield/models.py

@@ -1,3 +1,5 @@
+from typing import Any, Mapping, Optional, Sequence, Type, TypeVar, cast
+
 from django.db.models.fields import BigIntegerField
 
 from bitfield.query import BitQueryExactLookupStub
@@ -148,12 +150,21 @@ class BitField(BigIntegerField):
         return name, path, args, kwargs
 
 
+def flags_from_annotations(annotations: Mapping[str, type]) -> Sequence[str]:
+    flags = []
+    for attr, ty in annotations.items():
+        assert ty in ("bool", bool), f"bitfields can only hold bools, {attr} is {ty!r}"
+        flags.append(attr)
+
+    return flags
+
+
 class TypedBitfieldMeta(type):
     def __new__(cls, name, bases, clsdict):
-        if name == "TypedBitfield":
+        if name == "TypedClassBitField":
             return type.__new__(cls, name, bases, clsdict)
 
-        flags = []
+        flags = {}
         for attr, ty in clsdict["__annotations__"].items():
             if attr.startswith("_"):
                 continue
@@ -161,11 +172,10 @@ class TypedBitfieldMeta(type):
             if attr in ("bitfield_default", "bitfield_null"):
                 continue
 
-            assert ty in ("bool", bool), f"bitfields can only hold bools, {attr} is {ty!r}"
-            flags.append(attr)
+            flags[attr] = ty
 
         return BitField(
-            flags=flags,
+            flags=flags_from_annotations(flags),
             default=clsdict.get("bitfield_default"),
             null=clsdict.get("bitfield_null") or False,
         )
@@ -174,8 +184,37 @@ class TypedBitfieldMeta(type):
         raise NotImplementedError()
 
 
-class TypedBitfield(metaclass=TypedBitfieldMeta):
-    pass
+class TypedClassBitField(metaclass=TypedBitfieldMeta):
+    """
+    A wrapper around BitField that allows you to access its fields as instance
+    attributes in a type-safe way.
+    """
+
+    bitfield_default: Optional[Any]
+    bitfield_null: bool
+
+
+T = TypeVar("T")
+
+
+def typed_dict_bitfield(definition: Type[T], default=None, null=False) -> T:
+    """
+    A wrapper around BitField that allows you to access its fields as
+    dictionary keys attributes in a type-safe way.
+
+    Prefer `TypedClassBitField` over this if you can help it. This function
+    only exists to make it simpler to type bitfields with fields that are not
+    valid Python identifiers, but has limitations for how far it can provide
+    type safety.
+    """
+    assert issubclass(definition, dict)
+
+    return cast(
+        T,
+        BitField(
+            flags=flags_from_annotations(definition.__annotations__), default=default, null=null
+        ),
+    )
 
 
 BitField.register_lookup(BitQueryExactLookupStub)

src/sentry/migrations/0501_typed_bitfield_remove_labels.py

@@ -0,0 +1,176 @@
+# Generated by Django 2.2.28 on 2023-06-28 13:38
+
+from django.db import migrations
+
+import bitfield.models
+from sentry.new_migrations.migrations import CheckedMigration
+
+
+class Migration(CheckedMigration):
+    # This flag is used to mark that a migration shouldn't be automatically run in production. For
+    # the most part, this should only be used for operations where it's safe to run the migration
+    # after your code has deployed. So this should not be used for most operations that alter the
+    # schema of a table.
+    # Here are some things that make sense to mark as dangerous:
+    # - Large data migrations. Typically we want these to be run manually by ops so that they can
+    #   be monitored and not block the deploy for a long period of time while they run.
+    # - Adding indexes to large tables. Since this can take a long time, we'd generally prefer to
+    #   have ops run this and not block the deploy. Note that while adding an index is a schema
+    #   change, it's completely safe to run the operation after the code has deployed.
+    is_dangerous = False
+
+    dependencies = [
+        ("sentry", "0500_set_none_date_last_modified_to_date_uploaded"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="apiauthorization",
+            name="scopes",
+            field=bitfield.models.BitField(
+                [
+                    "project:read",
+                    "project:write",
+                    "project:admin",
+                    "project:releases",
+                    "team:read",
+                    "team:write",
+                    "team:admin",
+                    "event:read",
+                    "event:write",
+                    "event:admin",
+                    "org:read",
+                    "org:write",
+                    "org:admin",
+                    "member:read",
+                    "member:write",
+                    "member:admin",
+                ],
+                default=None,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="apigrant",
+            name="scopes",
+            field=bitfield.models.BitField(
+                [
+                    "project:read",
+                    "project:write",
+                    "project:admin",
+                    "project:releases",
+                    "team:read",
+                    "team:write",
+                    "team:admin",
+                    "event:read",
+                    "event:write",
+                    "event:admin",
+                    "org:read",
+                    "org:write",
+                    "org:admin",
+                    "member:read",
+                    "member:write",
+                    "member:admin",
+                ],
+                default=None,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="apikey",
+            name="scopes",
+            field=bitfield.models.BitField(
+                [
+                    "project:read",
+                    "project:write",
+                    "project:admin",
+                    "project:releases",
+                    "team:read",
+                    "team:write",
+                    "team:admin",
+                    "event:read",
+                    "event:write",
+                    "event:admin",
+                    "org:read",
+                    "org:write",
+                    "org:admin",
+                    "member:read",
+                    "member:write",
+                    "member:admin",
+                ],
+                default=None,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="apitoken",
+            name="scopes",
+            field=bitfield.models.BitField(
+                [
+                    "project:read",
+                    "project:write",
+                    "project:admin",
+                    "project:releases",
+                    "team:read",
+                    "team:write",
+                    "team:admin",
+                    "event:read",
+                    "event:write",
+                    "event:admin",
+                    "org:read",
+                    "org:write",
+                    "org:admin",
+                    "member:read",
+                    "member:write",
+                    "member:admin",
+                ],
+                default=None,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="authprovider",
+            name="flags",
+            field=bitfield.models.BitField(["allow_unlinked", "scim_enabled"], default=0),
+        ),
+        migrations.AlterField(
+            model_name="organizationmember",
+            name="flags",
+            field=bitfield.models.BitField(
+                [
+                    "sso:linked",
+                    "sso:invalid",
+                    "member-limit:restricted",
+                    "idp:provisioned",
+                    "idp:role-restricted",
+                ],
+                default=0,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="sentryapp",
+            name="scopes",
+            field=bitfield.models.BitField(
+                [
+                    "project:read",
+                    "project:write",
+                    "project:admin",
+                    "project:releases",
+                    "team:read",
+                    "team:write",
+                    "team:admin",
+                    "event:read",
+                    "event:write",
+                    "event:admin",
+                    "org:read",
+                    "org:write",
+                    "org:admin",
+                    "member:read",
+                    "member:write",
+                    "member:admin",
+                ],
+                default=None,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="user",
+            name="flags",
+            field=bitfield.models.BitField(["newsletter_consent_prompt"], default=0, null=True),
+        ),
+    ]

src/sentry/models/apigrant.py

@@ -1,10 +1,11 @@
 from datetime import timedelta
+from typing import TypedDict
 from uuid import uuid4
 
 from django.db import models
 from django.utils import timezone
 
-from bitfield import BitField
+from bitfield import typed_dict_bitfield
 from sentry.db.models import ArrayField, FlexibleForeignKey, Model, control_silo_only_model
 
 DEFAULT_EXPIRATION = timedelta(minutes=10)
@@ -33,24 +34,27 @@ class ApiGrant(Model):
     code = models.CharField(max_length=64, db_index=True, default=generate_code)
     expires_at = models.DateTimeField(db_index=True, default=default_expiration)
     redirect_uri = models.CharField(max_length=255)
-    scopes = BitField(
-        flags=(
-            ("project:read", "project:read"),
-            ("project:write", "project:write"),
-            ("project:admin", "project:admin"),
-            ("project:releases", "project:releases"),
-            ("team:read", "team:read"),
-            ("team:write", "team:write"),
-            ("team:admin", "team:admin"),
-            ("event:read", "event:read"),
-            ("event:write", "event:write"),
-            ("event:admin", "event:admin"),
-            ("org:read", "org:read"),
-            ("org:write", "org:write"),
-            ("org:admin", "org:admin"),
-            ("member:read", "member:read"),
-            ("member:write", "member:write"),
-            ("member:admin", "member:admin"),
+    scopes = typed_dict_bitfield(
+        TypedDict(  # type: ignore[operator]
+            "scopes",
+            {
+                "project:read": bool,
+                "project:write": bool,
+                "project:admin": bool,
+                "project:releases": bool,
+                "team:read": bool,
+                "team:write": bool,
+                "team:admin": bool,
+                "event:read": bool,
+                "event:write": bool,
+                "event:admin": bool,
+                "org:read": bool,
+                "org:write": bool,
+                "org:admin": bool,
+                "member:read": bool,
+                "member:write": bool,
+                "member:admin": bool,
+            },
         )
     )
     scope_list = ArrayField(of=models.TextField)

src/sentry/models/apikey.py

@@ -1,10 +1,11 @@
+from typing import TypedDict
 from uuid import uuid4
 
 from django.db import models
 from django.utils import timezone
 from django.utils.translation import ugettext_lazy as _
 
-from bitfield import BitField
+from bitfield import typed_dict_bitfield
 from sentry.db.models import (
     ArrayField,
     BaseManager,
@@ -29,24 +30,27 @@ class ApiKey(Model):
     organization_id = HybridCloudForeignKey("sentry.Organization", on_delete="cascade")
     label = models.CharField(max_length=64, blank=True, default="Default")
     key = models.CharField(max_length=32, unique=True)
-    scopes = BitField(
-        flags=(
-            ("project:read", "project:read"),
-            ("project:write", "project:write"),
-            ("project:admin", "project:admin"),
-            ("project:releases", "project:releases"),
-            ("team:read", "team:read"),
-            ("team:write", "team:write"),
-            ("team:admin", "team:admin"),
-            ("event:read", "event:read"),
-            ("event:write", "event:write"),
-            ("event:admin", "event:admin"),
-            ("org:read", "org:read"),
-            ("org:write", "org:write"),
-            ("org:admin", "org:admin"),
-            ("member:read", "member:read"),
-            ("member:write", "member:write"),
-            ("member:admin", "member:admin"),
+    scopes = typed_dict_bitfield(
+        TypedDict(  # type: ignore[operator]
+            "scopes",
+            {
+                "project:read": bool,
+                "project:write": bool,
+                "project:admin": bool,
+                "project:releases": bool,
+                "team:read": bool,
+                "team:write": bool,
+                "team:admin": bool,
+                "event:read": bool,
+                "event:write": bool,
+                "event:admin": bool,
+                "org:read": bool,
+                "org:write": bool,
+                "org:admin": bool,
+                "member:read": bool,
+                "member:write": bool,
+                "member:admin": bool,
+            },
         )
     )
     scope_list = ArrayField(of=models.TextField)

src/sentry/models/apiscopes.py

@@ -1,8 +1,9 @@
 from collections.abc import Sequence
+from typing import TypedDict
 
 from django.db import models
 
-from bitfield import BitField
+from bitfield import typed_dict_bitfield
 from sentry.db.models import ArrayField
 
 
@@ -26,9 +27,6 @@ class ApiScopes(Sequence):
             + self.__class__.member
         )
 
-    def to_bitfield(self):
-        return tuple((s, s) for s in self.scopes)
-
     def __getitem__(self, value):
         return self.scopes.__getitem__(value)
 
@@ -48,7 +46,29 @@ class HasApiScopes(models.Model):
         abstract = True
 
     # List of scopes in bit form
-    scopes = BitField(flags=ApiScopes().to_bitfield())
+    ScopesDict = TypedDict(
+        "ScopesDict",
+        {
+            "project:read": bool,
+            "project:write": bool,
+            "project:admin": bool,
+            "project:releases": bool,
+            "team:read": bool,
+            "team:write": bool,
+            "team:admin": bool,
+            "event:read": bool,
+            "event:write": bool,
+            "event:admin": bool,
+            "org:read": bool,
+            "org:write": bool,
+            "org:admin": bool,
+            "member:read": bool,
+            "member:write": bool,
+            "member:admin": bool,
+        },
+    )
+    assert set(ScopesDict.__annotations__) == set(ApiScopes())
+    scopes = typed_dict_bitfield(ScopesDict)
 
     # Human readable list of scopes
     scope_list = ArrayField(of=models.TextField)

src/sentry/models/authprovider.py

@@ -3,7 +3,7 @@ import logging
 from django.db import models
 from django.utils import timezone
 
-from bitfield import BitField
+from bitfield import TypedClassBitField
 from sentry.db.models import (
     BoundedBigIntegerField,
     BoundedPositiveIntegerField,
@@ -53,13 +53,13 @@ class AuthProvider(Model):
     default_role = BoundedPositiveIntegerField(default=50)
     default_global_access = models.BooleanField(default=True)
 
-    flags = BitField(
-        flags=(
-            ("allow_unlinked", "Grant access to members who have not linked SSO accounts."),
-            ("scim_enabled", "Enable SCIM for member and team provisioning and syncing"),
-        ),
-        default=0,
-    )
+    class flags(TypedClassBitField):
+        # Grant access to members who have not linked SSO accounts.
+        allow_unlinked: bool
+        # Enable SCIM for member and team provisioning and syncing.
+        scim_enabled: bool
+
+        bitfield_default = 0
 
     class Meta:
         app_label = "sentry"

src/sentry/models/organization.py

@@ -13,7 +13,7 @@ from django.utils import timezone
 from django.utils.functional import cached_property
 from typing_extensions import override
 
-from bitfield import TypedBitfield
+from bitfield import TypedClassBitField
 from sentry import features, roles
 from sentry.app import env
 from sentry.constants import (
@@ -86,7 +86,7 @@ OrganizationStatus._labels = {
 
 
 class OrganizationManager(BaseManager):
-    def get_for_user_ids(self, user_ids: Sequence[int]) -> QuerySet:
+    def get_for_user_ids(self, user_ids: Collection[int]) -> QuerySet:
         """Returns the QuerySet of all organizations that a set of Users have access to."""
         return self.filter(
             status=OrganizationStatus.ACTIVE,
@@ -175,7 +175,7 @@ class Organization(Model, OptionMixin, OrganizationAbsoluteUrlMixin, SnowflakeId
     default_role = models.CharField(max_length=32, default=str(roles.get_default().id))
     is_test = models.BooleanField(default=False)
 
-    class flags(TypedBitfield):
+    class flags(TypedClassBitField):
         # Allow members to join and leave teams without requiring approval
         allow_joinleave: bool