feat/entity types for bookmarks

alembic/versions/142f4198a2ca_migrate_bookmark_categories_to_entity_.py

@@ -0,0 +1,186 @@
+"""migrate_bookmark_categories_to_entity_types
+
+Revision ID: 142f4198a2ca
+Revises: 0af36743e9f3
+Create Date: 2025-08-18 14:32:23.555984
+
+"""
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "142f4198a2ca"
+down_revision: Union[str, None] = "0af36743e9f3"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Create new enum type with EntityType values
+    op.execute(
+        """
+        CREATE TYPE entitytype AS ENUM(
+            'analysis_software_source_code',
+            'brain_atlas',
+            'brain_atlas_region',
+            'cell_composition',
+            'electrical_cell_recording',
+            'electrical_recording_stimulus',
+            'emodel',
+            'experimental_bouton_density',
+            'experimental_neuron_density',
+            'experimental_synapses_per_connection',
+            'ion_channel_model',
+            'memodel',
+            'mesh',
+            'memodel_calibration_result',
+            'me_type_density',
+            'publication',
+            'reconstruction_morphology',
+            'simulation',
+            'simulation_campaign',
+            'simulation_campaign_generation',
+            'simulation_execution',
+            'simulation_result',
+            'scientific_artifact',
+            'single_neuron_simulation',
+            'single_neuron_synaptome',
+            'single_neuron_synaptome_simulation',
+            'subject',
+            'validation_result',
+            'circuit'
+        )
+    """
+    )
+
+    # Add temporary column with new enum type
+    op.execute("ALTER TABLE bookmark ADD COLUMN category_new entitytype")
+
+    # Migrate data from old enum to new enum using mapping
+    migrations = [
+        ("ExperimentalBoutonDensity", "experimental_bouton_density"),
+        ("ExperimentalNeuronDensity", "experimental_neuron_density"),
+        ("ExperimentalElectroPhysiology", "electrical_cell_recording"),
+        ("ExperimentalSynapsePerConnection", "experimental_synapses_per_connection"),
+        ("ExperimentalNeuronMorphology", "reconstruction_morphology"),
+        ("SimulationCampaign", "simulation_campaign"),
+        ("CircuitEModel", "emodel"),
+        ("CircuitMEModel", "memodel"),
+        ("SingleNeuronSynaptome", "single_neuron_synaptome"),
+        ("SingleNeuronSimulation", "single_neuron_simulation"),
+        ("SynaptomeSimulation", "single_neuron_synaptome_simulation"),
+    ]
+
+    for old_value, new_value in migrations:
+        op.execute(
+            f"""
+            UPDATE bookmark 
+            SET category_new = '{new_value}' 
+            WHERE category = '{old_value}'
+        """
+        )
+
+    # Drop old unique constraint (resource_id based)
+    op.drop_constraint(
+        "bookmark_unique_for_resource_category_per_project", "bookmark", type_="unique"
+    )
+
+    # Drop resource_id column (no longer used in new model)
+    op.drop_index("ix_bookmark_resource_id", table_name="bookmark")
+    op.drop_column("bookmark", "resource_id")
+
+    # Drop old column and rename new column
+    op.execute("ALTER TABLE bookmark DROP COLUMN category")
+    op.execute("ALTER TABLE bookmark RENAME COLUMN category_new TO category")
+
+    # Set NOT NULL constraint on category column
+    op.execute("ALTER TABLE bookmark ALTER COLUMN category SET NOT NULL")
+
+    # Create new unique constraint (entity_id based)
+    op.create_unique_constraint(
+        "bookmark_unique_for_entity_category_per_project",
+        "bookmark",
+        ["entity_id", "category", "project_id"],
+    )
+
+    # Drop old enum type
+    op.execute("DROP TYPE bookmarkcategory")
+
+
+def downgrade() -> None:
+    # Recreate old enum type
+    op.execute(
+        """
+        CREATE TYPE bookmarkcategory AS ENUM(
+            'ExperimentalBoutonDensity',
+            'ExperimentalNeuronDensity', 
+            'ExperimentalElectroPhysiology',
+            'ExperimentalSynapsePerConnection',
+            'ExperimentalNeuronMorphology',
+            'SimulationCampaign',
+            'CircuitEModel',
+            'CircuitMEModel',
+            'SingleNeuronSynaptome',
+            'SingleNeuronSimulation',
+            'SynaptomeSimulation'
+        )
+    """
+    )
+
+    # Add temporary column with old enum type
+    op.execute("ALTER TABLE bookmark ADD COLUMN category_old bookmarkcategory")
+
+    # Migrate data back from new enum to old enum
+    reverse_migrations = [
+        ("experimental_bouton_density", "ExperimentalBoutonDensity"),
+        ("experimental_neuron_density", "ExperimentalNeuronDensity"),
+        ("electrical_cell_recording", "ExperimentalElectroPhysiology"),
+        ("experimental_synapses_per_connection", "ExperimentalSynapsePerConnection"),
+        ("reconstruction_morphology", "ExperimentalNeuronMorphology"),
+        ("simulation_campaign", "SimulationCampaign"),
+        ("emodel", "CircuitEModel"),
+        ("memodel", "CircuitMEModel"),
+        ("single_neuron_synaptome", "SingleNeuronSynaptome"),
+        ("single_neuron_simulation", "SingleNeuronSimulation"),
+        ("single_neuron_synaptome_simulation", "SynaptomeSimulation"),
+    ]
+
+    for new_value, old_value in reverse_migrations:
+        op.execute(
+            f"""
+            UPDATE bookmark 
+            SET category_old = '{old_value}' 
+            WHERE category = '{new_value}'
+        """
+        )
+
+    # Drop new unique constraint (entity_id based)
+    op.drop_constraint(
+        "bookmark_unique_for_entity_category_per_project", "bookmark", type_="unique"
+    )
+
+    # Add back resource_id column (for old model compatibility)
+    op.add_column("bookmark", sa.Column("resource_id", sa.String(), nullable=True))
+    op.create_index(
+        "ix_bookmark_resource_id", "bookmark", ["resource_id"], unique=False
+    )
+
+    # Drop new column and rename old column back
+    op.execute("ALTER TABLE bookmark DROP COLUMN category")
+    op.execute("ALTER TABLE bookmark RENAME COLUMN category_old TO category")
+
+    # Set NOT NULL constraint on category column
+    op.execute("ALTER TABLE bookmark ALTER COLUMN category SET NOT NULL")
+
+    # Recreate old unique constraint (resource_id based)
+    op.create_unique_constraint(
+        "bookmark_unique_for_resource_category_per_project",
+        "bookmark",
+        ["resource_id", "category", "project_id"],
+    )
+
+    # Drop new enum type
+    op.execute("DROP TYPE entitytype")

virtual_labs/domain/bookmark.py

@@ -1,27 +1,46 @@
-from enum import Enum
-from typing import Annotated, TypedDict
+from enum import StrEnum, auto
+from typing import TypedDict
 
-from pydantic import UUID4, BaseModel, Field
+from pydantic import UUID4, BaseModel
 
 
-class BookmarkCategory(Enum):
-    ExperimentalBoutonDensity = "ExperimentalBoutonDensity"
-    ExperimentalNeuronDensity = "ExperimentalNeuronDensity"
-    ExperimentalElectroPhysiology = "ExperimentalElectroPhysiology"
-    ExperimentalSynapsePerConnection = "ExperimentalSynapsePerConnection"
-    ExperimentalNeuronMorphology = "ExperimentalNeuronMorphology"
-    SimulationCampaign = "SimulationCampaign"
-    CircuitEModel = "CircuitEModel"
-    CircuitMEModel = "CircuitMEModel"
-    SingleNeuronSynaptome = "SingleNeuronSynaptome"
-    SingleNeuronSimulation = "SingleNeuronSimulation"
-    SynaptomeSimulation = "SynaptomeSimulation"
+class EntityType(StrEnum):
+    """Entity types."""
+
+    analysis_software_source_code = auto()
+    brain_atlas = auto()
+    brain_atlas_region = auto()
+    cell_composition = auto()
+    electrical_cell_recording = auto()
+    electrical_recording_stimulus = auto()
+    emodel = auto()
+    experimental_bouton_density = auto()
+    experimental_neuron_density = auto()
+    experimental_synapses_per_connection = auto()
+    ion_channel_model = auto()
+    memodel = auto()
+    mesh = auto()
+    memodel_calibration_result = auto()
+    me_type_density = auto()
+    publication = auto()
+    reconstruction_morphology = auto()
+    simulation = auto()
+    simulation_campaign = auto()
+    simulation_campaign_generation = auto()
+    simulation_execution = auto()
+    simulation_result = auto()
+    scientific_artifact = auto()
+    single_neuron_simulation = auto()
+    single_neuron_synaptome = auto()
+    single_neuron_synaptome_simulation = auto()
+    subject = auto()
+    validation_result = auto()
+    circuit = auto()
 
 
 class BookmarkIn(BaseModel):
-    entity_id: UUID4 | None = None
-    resource_id: Annotated[str | None, Field(alias="resourceId")] = None
-    category: BookmarkCategory
+    entity_id: UUID4
+    category: EntityType
 
     class Config:
         from_attributes = True
@@ -31,9 +50,7 @@ class Config:
 class BookmarkOut(BaseModel):
     id: UUID4
     entity_id: UUID4 | None = None
-    # TODO: return snake case for entity core
-    resource_id: Annotated[str | None, Field(alias="resourceId")] = None
-    category: BookmarkCategory
+    category: EntityType
 
     class Config:
         from_attributes = True

virtual_labs/infrastructure/db/models.py

@@ -25,7 +25,7 @@
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
 from sqlalchemy.sql import func
 
-from virtual_labs.domain.bookmark import BookmarkCategory
+from virtual_labs.domain.bookmark import EntityType
 
 
 class Base(DeclarativeBase):
@@ -211,11 +211,10 @@ class Bookmark(Base):
     __tablename__ = "bookmark"
 
     id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
-    resource_id = Column(String, nullable=True, index=True)
     entity_id: Mapped[uuid.UUID | None] = mapped_column(
         unique=False, index=True, nullable=True
     )
-    category = Column(SAEnum(BookmarkCategory), nullable=False)  # type: ignore[var-annotated]
+    category: Mapped[EntityType] = mapped_column(SAEnum(EntityType), nullable=False)
 
     created_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True), default=func.now()
@@ -226,10 +225,10 @@ class Bookmark(Base):
 
     __table_args__ = (
         UniqueConstraint(
-            resource_id,
+            entity_id,
             category,
             project_id,
-            name="bookmark_unique_for_resource_category_per_project",
+            name="bookmark_unique_for_entity_category_per_project",
         ),
     )