diff --git a/pyproject.toml b/pyproject.toml index 392e1bd40067cd..117ec9873dc041 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -446,7 +446,6 @@ module = [ "sentry.utils.auth", "sentry.utils.committers", "sentry.utils.services", - "sentry.utils.snowflake", "sentry.web.forms.accounts", "sentry.web.frontend.account_identity", "sentry.web.frontend.auth_close", @@ -643,6 +642,7 @@ module = [ "sentry.utils.redis_metrics", "sentry.utils.sentry_apps.*", "sentry.utils.sms", + "sentry.utils.snowflake", "sentry.utils.urls", "sentry.utils.uwsgi", "sentry.utils.zip", diff --git a/src/sentry/utils/snowflake.py b/src/sentry/utils/snowflake.py index 594739e197f5ee..de0b6a6cbd6d8f 100644 --- a/src/sentry/utils/snowflake.py +++ b/src/sentry/utils/snowflake.py @@ -1,16 +1,22 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta +from typing import TYPE_CHECKING from django.conf import settings -from django.db import IntegrityError, models, router, transaction +from django.db import IntegrityError, router, transaction +from redis.client import StrictRedis from rest_framework import status from rest_framework.exceptions import APIException from sentry.db.postgres.transactions import enforce_constraints from sentry.types.region import RegionContextError, get_local_region +if TYPE_CHECKING: + from sentry.db.models.base import Model + _TTL = timedelta(minutes=5) @@ -19,19 +25,31 @@ class MaxSnowflakeRetryError(APIException): default_detail = "Max allowed ID retry reached. Please try again in a second" -class SnowflakeIdMixin: - id: models.Field[int, int] +def _snowflake_inst_is_model(inst: object) -> Model: + # this is an unsound hack to make the mixin typing work + # ideally the mixin should just be a function + from sentry.db.models.base import Model + + if not isinstance(inst, Model): + raise TypeError(f"expected SnowflakeIdMixin to be mixed into Model, got {type(inst)}") + else: + return inst - def save_with_snowflake_id(self, snowflake_redis_key, save_callback): + +class SnowflakeIdMixin: + def save_with_snowflake_id( + self, snowflake_redis_key: str, save_callback: Callable[[], object] + ) -> None: + inst = _snowflake_inst_is_model(self) for _ in range(settings.MAX_REDIS_SNOWFLAKE_RETRY_COUNTER): - if not self.id: - self.id = generate_snowflake_id(snowflake_redis_key) + if not inst.id: + inst.id = generate_snowflake_id(snowflake_redis_key) try: - with enforce_constraints(transaction.atomic(using=router.db_for_write(type(self)))): + with enforce_constraints(transaction.atomic(using=router.db_for_write(type(inst)))): save_callback() return except IntegrityError: - self.id = None + inst.id = None # type: ignore[assignment] # see typeddjango/django-stubs#2014 raise MaxSnowflakeRetryError @@ -40,11 +58,11 @@ class SnowflakeBitSegment: length: int name: str - def __post_init__(self): + def __post_init__(self) -> None: if self.length <= 0: raise Exception("The length should be a positive number") - def validate(self, value): + def validate(self, value: int) -> None: if value >> self.length != 0: raise Exception(f"{self.name} exceed max bit value of {self.length}") @@ -67,7 +85,7 @@ def validate(self, value): NULL_REGION_ID = 0 -def msb_0_ordering(value, width): +def msb_0_ordering(value: int, width: int) -> int: """ MSB 0 Ordering refers to when the bit numbering starts at zero for the most significant bit (MSB) the numbering scheme is called MSB 0. @@ -106,7 +124,7 @@ def generate_snowflake_id(redis_key: str) -> int: return snowflake_id -def get_redis_cluster(redis_key: str): +def get_redis_cluster(redis_key: str) -> StrictRedis[str]: from sentry.utils import redis return redis.clusters.get("default").get_local_client_for_key(redis_key)