diff --git a/surfsense_backend/alembic/versions/16_fix_connector_unique_constraint.py b/surfsense_backend/alembic/versions/16_fix_connector_unique_constraint.py new file mode 100644 index 0000000..38c84d3 --- /dev/null +++ b/surfsense_backend/alembic/versions/16_fix_connector_unique_constraint.py @@ -0,0 +1,76 @@ +"""Fix SearchSourceConnector unique constraint to allow per-user connectors + +Revision ID: '16' +Revises: '15' +Create Date: 2025-01-03 12:00:00.000000 + +""" + +from collections.abc import Sequence + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "16" +down_revision: str | None = "15" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Drop unique constraint on connector_type and add composite unique constraint on (user_id, connector_type).""" + + # First, drop the existing unique constraint on connector_type + # Note: PostgreSQL auto-generates constraint names, so we need to find and drop it + op.execute( + """ + DO $$ + DECLARE + constraint_name TEXT; + BEGIN + -- Find the unique constraint on connector_type column + SELECT tc.constraint_name INTO constraint_name + FROM information_schema.table_constraints tc + JOIN information_schema.key_column_usage kcu + ON tc.constraint_name = kcu.constraint_name + AND tc.table_schema = kcu.table_schema + WHERE tc.constraint_type = 'UNIQUE' + AND tc.table_name = 'search_source_connectors' + AND kcu.column_name = 'connector_type' + AND tc.table_schema = 'public'; + + -- Drop the constraint if it exists + IF constraint_name IS NOT NULL THEN + EXECUTE format('ALTER TABLE search_source_connectors DROP CONSTRAINT %I', constraint_name); + RAISE NOTICE 'Dropped unique constraint: %', constraint_name; + ELSE + RAISE NOTICE 'No unique constraint found on connector_type column'; + END IF; + END $$; + """ + ) + + # Add the new composite unique constraint + op.create_unique_constraint( + "uq_user_connector_type", + "search_source_connectors", + ["user_id", "connector_type"], + ) + + +def downgrade() -> None: + """Revert to unique constraint on connector_type only.""" + + # Drop the composite unique constraint + op.drop_constraint( + "uq_user_connector_type", "search_source_connectors", type_="unique" + ) + + # Add back the unique constraint on connector_type + # Note: This downgrade will fail if there are duplicate connector_types for different users + # In that case, manual cleanup would be required + op.create_unique_constraint( + None, # Let PostgreSQL auto-generate the constraint name + "search_source_connectors", + ["connector_type"], + ) diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index 1b7da7e..d749b3b 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -16,6 +16,7 @@ from sqlalchemy import ( Integer, String, Text, + UniqueConstraint, text, ) from sqlalchemy.dialects.postgresql import UUID @@ -227,11 +228,12 @@ class SearchSpace(BaseModel, TimestampMixin): class SearchSourceConnector(BaseModel, TimestampMixin): __tablename__ = "search_source_connectors" + __table_args__ = ( + UniqueConstraint("user_id", "connector_type", name="uq_user_connector_type"), + ) name = Column(String(100), nullable=False, index=True) - connector_type = Column( - SQLAlchemyEnum(SearchSourceConnectorType), nullable=False, unique=True - ) + connector_type = Column(SQLAlchemyEnum(SearchSourceConnectorType), nullable=False) is_indexable = Column(Boolean, nullable=False, default=False) last_indexed_at = Column(TIMESTAMP(timezone=True), nullable=True) config = Column(JSON, nullable=False)