feat: add missing FK indexes, unique constraints, and cleanup orphaned DAG tables (airweave-ai#1506)

* feat: integrate MyPy type checking into CI and pre-commit

Adds a dedicated MyPy job to the GitHub Actions workflow, running type checks on changed Python files within the `backend/airweave` directory. This optimizes execution time by focusing only on relevant modifications.

Introduces a pre-commit hook for MyPy, ensuring type correctness locally before commits are pushed.

Incorporates a `ruff check` step into the `code-quality.yml` workflow for explicit linting, complementing the existing formatting check.

* fix: point at /backend

* ci: scope ruff lint and format checks to changed files only

* feat: add missing FK indexes, unique constraints, and pre-flight check

Adds indexes on unindexed foreign keys (sync_job.sync_id, sync_connection,
source_connection, connection, collection, entity_relation) and utility
indexes (organization.auth0_org_id, expires_at on session tables).

Adds unique constraints on sync_connection(sync_id, connection_id) and
entity_relation(from, to, name, org_id) to prevent duplicate rows.

Removes redundant idx_entity_count_sync_def (duplicate of unique constraint).

Includes pre-flight check script to validate no duplicate data exists
before applying the migration.

* feat: add alembic migration for indexes, constraints, and DAG table cleanup

Autogenerated migration that:
- Creates missing FK indexes on 7 tables
- Adds unique constraints on sync_connection and entity_relation
- Drops redundant idx_entity_count_sync_def
- Adds utility indexes (auth0_org_id, expires_at)
- Drops orphaned DAG tables (sync_dag, dag_node, dag_edge)
- Fixes access_control_membership FK to include ondelete CASCADE

* feat: drop useless B-tree index on search_queries.query_text

B-tree on unbounded Text is only useful for exact match or prefix.
Removing it saves write overhead on every INSERT into search_queries.

* Delete backend/scripts/preflight_check_model_indexes.py

Author: orhanrauf

backend/airweave/models/collection.py

@@ -1,9 +1,9 @@
 """Collection model."""
 
 import uuid
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Optional
 
-from sqlalchemy import ForeignKey, String
+from sqlalchemy import ForeignKey, Index, String
 from sqlalchemy.dialects.postgresql import JSONB, UUID
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
@@ -32,10 +32,6 @@ class Collection(OrganizationBase, UserMixin):
         "VectorDbDeploymentMetadata", lazy="joined"
     )
 
-    if TYPE_CHECKING:
-        search_queries: List["SearchQuery"]
-        source_connections: List["SourceConnection"]
-
     source_connections: Mapped[list["SourceConnection"]] = relationship(
         "SourceConnection",
         back_populates="collection",
@@ -51,3 +47,5 @@ class Collection(OrganizationBase, UserMixin):
         cascade="all, delete-orphan",
         passive_deletes=True,
     )
+
+    __table_args__ = (Index("idx_collection_vdb_metadata_id", "vector_db_deployment_metadata_id"),)

backend/airweave/models/connection.py

@@ -1,9 +1,9 @@
 """Connection model."""
 
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
-from sqlalchemy import CheckConstraint, ForeignKey, String, Text, event
+from sqlalchemy import CheckConstraint, ForeignKey, Index, String, Text, event
 from sqlalchemy import Enum as SQLAlchemyEnum
 from sqlalchemy.orm import Mapped, Session, mapped_column, relationship
 
@@ -115,8 +115,6 @@ class Connection(Base):
     )
 
     __table_args__ = (
-        # Enforce that organization_id, created_by_email, and modified_by_email are not null
-        # except for the specific native connections
         CheckConstraint(
             """
             (short_name IN ('qdrant_native', 'neo4j_native', 'local_text2vec'))
@@ -127,12 +125,14 @@ class Connection(Base):
             """,
             name="ck_connection_native_or_complete",
         ),
+        Index("idx_connection_organization_id", "organization_id"),
+        Index("idx_connection_integration_credential_id", "integration_credential_id"),
     )
 
 
 # Event to delete integration credential when Connection is deleted
 @event.listens_for(Connection, "before_delete")
-def delete_integration_credential(mapper, connection, target):
+def delete_integration_credential(mapper: Any, connection: Any, target: Any) -> None:
     """When a Connection is deleted, also delete its IntegrationCredential if present."""
     if target.integration_credential_id:
         # Get the session

backend/airweave/models/connection_init_session.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Optional
 from uuid import UUID
 
-from sqlalchemy import JSON, DateTime, ForeignKey, String
+from sqlalchemy import JSON, DateTime, ForeignKey, Index, String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from airweave.models._base import OrganizationBase
@@ -70,6 +70,8 @@ class ConnectionInitSession(OrganizationBase):
 
     redirect_session: Mapped[Optional["RedirectSession"]] = relationship("RedirectSession")
 
+    __table_args__ = (Index("idx_connection_init_session_expires_at", "expires_at"),)
+
     @staticmethod
     def default_expires_at(minutes: int = 30) -> datetime:
         """Return a UTC expiry timestamp ``minutes`` from now."""

backend/airweave/models/entity_count.py

@@ -53,5 +53,4 @@ class EntityCount(Base):
         ),
         Index("idx_entity_count_sync_id", "sync_id"),
         Index("idx_entity_count_entity_def_id", "entity_definition_id"),
-        Index("idx_entity_count_sync_def", "sync_id", "entity_definition_id"),
     )

backend/airweave/models/entity_relation.py

@@ -1,6 +1,10 @@
 """Models for entity relations."""
 
-from sqlalchemy import Column, ForeignKey, String
+from typing import Optional
+from uuid import UUID
+
+from sqlalchemy import ForeignKey, Index, String, UniqueConstraint
+from sqlalchemy.orm import Mapped, mapped_column
 
 from airweave.models._base import Base
 
@@ -10,8 +14,26 @@ class EntityRelation(Base):
 
     __tablename__ = "entity_relation"
 
-    name = Column(String, nullable=False)
-    description = Column(String)
-    from_entity_definition_id = Column(ForeignKey("entity_definition.id"), nullable=False)
-    to_entity_definition_id = Column(ForeignKey("entity_definition.id"), nullable=False)
-    organization_id = Column(ForeignKey("organization.id"), nullable=True)
+    name: Mapped[str] = mapped_column(String, nullable=False)
+    description: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+    from_entity_definition_id: Mapped[UUID] = mapped_column(
+        ForeignKey("entity_definition.id"), nullable=False
+    )
+    to_entity_definition_id: Mapped[UUID] = mapped_column(
+        ForeignKey("entity_definition.id"), nullable=False
+    )
+    organization_id: Mapped[Optional[UUID]] = mapped_column(
+        ForeignKey("organization.id"), nullable=True
+    )
+
+    __table_args__ = (
+        UniqueConstraint(
+            "from_entity_definition_id",
+            "to_entity_definition_id",
+            "name",
+            "organization_id",
+            name="uq_entity_relation",
+        ),
+        Index("idx_entity_relation_from", "from_entity_definition_id"),
+        Index("idx_entity_relation_to", "to_entity_definition_id"),
+    )

backend/airweave/models/organization.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING, List, Optional
 
-from sqlalchemy import JSON, String, Text
+from sqlalchemy import JSON, Index, String, Text
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from airweave.models._base import Base
@@ -77,3 +77,5 @@ class Organization(Base):
         cascade="all, delete-orphan",
         lazy="noload",
     )
+
+    __table_args__ = (Index("idx_organization_auth0_org_id", "auth0_org_id"),)

backend/airweave/models/organization_billing.py

@@ -17,7 +17,8 @@ class OrganizationBilling(Base):
 
     __tablename__ = "organization_billing"
 
-    # Foreign key to organization
+    # TODO: organization_id is typed as str but organization.id is UUID.
+    # Needs a data migration to fix — do not change the type without one.
     organization_id: Mapped[str] = mapped_column(
         ForeignKey("organization.id", ondelete="CASCADE"), unique=True
     )

backend/airweave/models/redirect_session.py

@@ -2,7 +2,7 @@
 
 from datetime import datetime, timedelta, timezone
 
-from sqlalchemy import DateTime, String, Text
+from sqlalchemy import DateTime, Index, String, Text
 from sqlalchemy.orm import Mapped, mapped_column
 
 from airweave.models._base import OrganizationBase
@@ -29,6 +29,8 @@ class RedirectSession(OrganizationBase):
     # Expiry; entries are consumed (deleted) on first use or when expired
     expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
 
+    __table_args__ = (Index("idx_redirect_session_expires_at", "expires_at"),)
+
     @staticmethod
     def default_expires_at(minutes: int = 5) -> datetime:
         """Generate a default expiration datetime for redirect sessions.

backend/airweave/models/search_query.py

@@ -124,7 +124,6 @@ class SearchQuery(OrganizationBase, UserMixin):
         Index("ix_search_queries_api_key_created", "api_key_id", "created_at"),
         Index("ix_search_queries_is_streaming", "is_streaming"),
         Index("ix_search_queries_retrieval_strategy", "retrieval_strategy"),
-        Index("ix_search_queries_query_text", "query_text"),  # For text analysis
         Index("ix_search_queries_duration", "duration_ms"),
         Index("ix_search_queries_results_count", "results_count"),
     )

backend/airweave/models/source_connection.py

@@ -1,10 +1,10 @@
 """Source connection model."""
 
 from time import sleep
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 from uuid import UUID
 
-from sqlalchemy import JSON, Boolean, ForeignKey, String, Text, event
+from sqlalchemy import JSON, Boolean, ForeignKey, Index, String, Text, event
 from sqlalchemy.orm import Mapped, Session, mapped_column, relationship
 
 from airweave.models._base import OrganizationBase, UserMixin
@@ -90,10 +90,16 @@ class SourceConnection(OrganizationBase, UserMixin):
         lazy="noload",
     )
 
+    __table_args__ = (
+        Index("idx_source_connection_sync_id", "sync_id"),
+        Index("idx_source_connection_connection_id", "connection_id"),
+        Index("idx_source_connection_collection_id", "readable_collection_id"),
+    )
+
 
 # Event to delete parent Sync when SourceConnection is deleted
 @event.listens_for(SourceConnection, "before_delete")
-def delete_parent_sync_and_connection(mapper, connection, target):
+def delete_parent_sync_and_connection(mapper: Any, connection: Any, target: Any) -> None:
     """When a SourceConnection is deleted, also delete its parent Sync and Connection."""
     # Delete parent Sync if it exists
     if target.sync_id: