typing and other small fixes

Author: wedamija

src/sentry/new_migrations/monkey/models.py

@@ -1,17 +1,16 @@
 from django.db.migrations import DeleteModel
 from django_zero_downtime_migrations.backends.postgres.schema import UnsafeOperationException
 
-from sentry.new_migrations.monkey.state import DeletionAction
+from sentry.db.postgres.schema import SafePostgresDatabaseSchemaEditor
+from sentry.new_migrations.monkey.state import DeletionAction, SentryProjectState
 
 
 class SafeDeleteModel(DeleteModel):
-    def __init__(self, *args, deletion_action: DeletionAction | None = None, **kwargs):
-        if deletion_action is None:
-            raise UnsafeOperationException("Deletion State is required")
+    def __init__(self, *args, deletion_action: DeletionAction, **kwargs):
         super().__init__(*args, **kwargs)
         self.deletion_action = deletion_action
 
-    def state_forwards(self, app_label, state):
+    def state_forwards(self, app_label: str, state: SentryProjectState) -> None:  # type: ignore[override]
         if self.deletion_action == DeletionAction.MOVE_TO_PENDING:
             model = state.apps.get_model(app_label, self.name)
             fields_with_constraints = [
@@ -25,22 +24,34 @@ def state_forwards(self, app_label, state):
                 )
         state.remove_model(app_label, self.name_lower, deletion_action=self.deletion_action)
 
-    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+    def database_forwards(
+        self,
+        app_label: str,
+        schema_editor: SafePostgresDatabaseSchemaEditor,  # type: ignore[override]
+        from_state: SentryProjectState,  # type: ignore[override]
+        to_state: SentryProjectState,  # type: ignore[override]
+    ) -> None:
         if self.deletion_action == DeletionAction.MOVE_TO_PENDING:
             return
 
         model = from_state.get_pending_deletion_model(app_label, self.name)
         if self.allow_migrate_model(schema_editor.connection.alias, model):
             schema_editor.delete_model(model, is_safe=True)
 
-    def database_backwards(self, app_label, schema_editor, from_state, to_state):
+    def database_backwards(
+        self,
+        app_label: str,
+        schema_editor: SafePostgresDatabaseSchemaEditor,  # type: ignore[override]
+        from_state: SentryProjectState,  # type: ignore[override]
+        to_state: SentryProjectState,  # type: ignore[override]
+    ) -> None:
         if self.deletion_action == DeletionAction.MOVE_TO_PENDING:
             return
         model = to_state.get_pending_deletion_model(app_label, self.name)
         if self.allow_migrate_model(schema_editor.connection.alias, model):
             schema_editor.create_model(model)
 
-    def describe(self):
+    def describe(self) -> str:
         if self.deletion_action == DeletionAction.MOVE_TO_PENDING:
             return f"Moved model {self.name} to pending deletion state"
         else:

src/sentry/new_migrations/monkey/state.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 from copy import deepcopy
 from enum import Enum
 
 from django.db.migrations.state import ProjectState
+from django.db.models import Model
 from django_zero_downtime_migrations.backends.postgres.schema import UnsafeOperationException
 
 
@@ -13,9 +16,9 @@ class DeletionAction(Enum):
 class SentryProjectState(ProjectState):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.pending_deletion_models = {}
+        self.pending_deletion_models: dict[tuple[str, str], type[Model]] = {}
 
-    def get_pending_deletion_model(self, app_label, model_name):
+    def get_pending_deletion_model(self, app_label: str, model_name: str) -> type[Model]:
         model_key = (app_label.lower(), model_name.lower())
         if model_key not in self.pending_deletion_models:
             raise UnsafeOperationException(
@@ -24,8 +27,10 @@ def get_pending_deletion_model(self, app_label, model_name):
             )
         return self.pending_deletion_models[model_key]
 
-    def remove_model(self, app_label, model_name, deletion_action: DeletionAction | None = None):
-        model_key = (app_label, model_name)
+    def remove_model(
+        self, app_label: str, model_name: str, deletion_action: DeletionAction | None = None
+    ) -> None:
+        model_key = (app_label.lower(), model_name.lower())
         if deletion_action == DeletionAction.DELETE:
             if model_key not in self.pending_deletion_models:
                 raise UnsafeOperationException(
@@ -43,7 +48,7 @@ def remove_model(self, app_label, model_name, deletion_action: DeletionAction |
             self.pending_deletion_models[model_key] = self.apps.get_model(app_label, model_name)
         super().remove_model(app_label, model_name)
 
-    def clone(self):
+    def clone(self) -> SentryProjectState:
         new_state = super().clone()
         new_state.pending_deletion_models = deepcopy(self.pending_deletion_models)  # type: ignore[attr-defined]
-        return new_state
+        return new_state  # type: ignore[return-value]