server/custom_field: delete association table records when deleting a…

… custom field

server/polar/custom_field/attachment.py

@@ -1,9 +1,15 @@
-from typing import TYPE_CHECKING, Annotated
+from typing import TYPE_CHECKING, Annotated, Any
 from uuid import UUID
 
 from pydantic import UUID4, Field
-from sqlalchemy import Boolean, ForeignKey, Integer, Uuid
-from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship
+from sqlalchemy import Boolean, ForeignKey, Integer, Uuid, event
+from sqlalchemy.orm import (
+    Mapped,
+    Mapper,
+    declared_attr,
+    mapped_column,
+    relationship,
+)
 
 from polar.kit.schemas import Schema
 
@@ -30,6 +36,16 @@ def custom_field(cls) -> Mapped["CustomField"]:
         return relationship("CustomField", lazy="joined")
 
 
+attached_custom_fields_models: set[type[AttachedCustomFieldMixin]] = set()
+
+
+# Event listener to track models inheriting from AttachedCustomFieldMixin
+@event.listens_for(Mapper, "mapper_configured")
+def track_attached_custom_field_mixin(_mapper: Mapper[Any], class_: type) -> None:
+    if issubclass(class_, AttachedCustomFieldMixin):
+        attached_custom_fields_models.add(class_)
+
+
 class AttachedCustomField(Schema):
     """Schema of a custom field attached to a resource."""
 

server/polar/custom_field/service.py

@@ -2,7 +2,7 @@
 from collections.abc import Sequence
 from typing import Any
 
-from sqlalchemy import Select, UnaryExpression, asc, desc, or_, select
+from sqlalchemy import Select, UnaryExpression, asc, delete, desc, or_, select
 
 from polar.auth.models import AuthSubject, is_organization, is_user
 from polar.authz.service import AccessType, Authz
@@ -16,6 +16,7 @@
 from polar.organization.resolver import get_payload_organization
 from polar.postgres import AsyncSession
 
+from .attachment import attached_custom_fields_models
 from .schemas import CustomFieldCreate, CustomFieldUpdate
 
 
@@ -165,6 +166,14 @@ async def delete(
     ) -> CustomField:
         custom_field.set_deleted_at()
         session.add(custom_field)
+
+        # Delete row with this custom field from all association tables
+        for model in attached_custom_fields_models:
+            delete_statement = delete(model).where(
+                model.custom_field_id == custom_field.id
+            )
+            await session.execute(delete_statement)
+
         return custom_field
 
     async def get_by_organization_and_id(