feat: dataset labels migration (#9409)

Author: ArcticFaded

scripts/ddl/postgresql_schema.sql

@@ -509,6 +509,48 @@ CREATE INDEX ix_api_keys_user_id ON public.api_keys
     USING btree (user_id);
 
 
+-- Table: dataset_labels
+-- ---------------------
+CREATE TABLE public.dataset_labels (
+    id serial NOT NULL,
+    name VARCHAR NOT NULL,
+    description VARCHAR,
+    color VARCHAR NOT NULL DEFAULT '#ffffff'::character varying,
+    user_id INTEGER,
+    CONSTRAINT pk_dataset_labels PRIMARY KEY (id),
+    CONSTRAINT uq_dataset_labels_name
+        UNIQUE (name),
+    CONSTRAINT fk_dataset_labels_user_id_users FOREIGN KEY
+        (user_id)
+        REFERENCES public.users (id)
+        ON DELETE SET NULL
+);
+
+
+-- Table: datasets_dataset_labels
+-- ------------------------------
+CREATE TABLE public.datasets_dataset_labels (
+    id serial NOT NULL,
+    dataset_id INTEGER NOT NULL,
+    dataset_label_id INTEGER NOT NULL,
+    CONSTRAINT pk_datasets_dataset_labels PRIMARY KEY (id),
+    CONSTRAINT uq_datasets_dataset_labels_dataset_id_dataset_label_id
+        UNIQUE (dataset_id, dataset_label_id),
+    CONSTRAINT fk_datasets_dataset_labels_dataset_id_datasets FOREIGN KEY
+        (dataset_id)
+        REFERENCES public.datasets (id)
+        ON DELETE CASCADE,
+    CONSTRAINT fk_datasets_dataset_labels_dataset_label_id_dataset_labels
+        FOREIGN KEY
+        (dataset_label_id)
+        REFERENCES public.dataset_labels (id)
+        ON DELETE CASCADE
+);
+
+CREATE INDEX ix_datasets_dataset_labels_dataset_label_id ON public.datasets_dataset_labels
+    USING btree (dataset_label_id);
+
+
 -- Table: document_annotations
 -- ---------------------------
 CREATE TABLE public.document_annotations (

src/phoenix/db/migrations/versions/58228d933c91_dataset_labels.py

@@ -0,0 +1,58 @@
+"""dataset_labels
+
+Revision ID: 58228d933c91
+Revises: 699f655af132
+Create Date: 2025-09-05 17:47:34.637329
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "58228d933c91"
+down_revision: Union[str, None] = "699f655af132"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "dataset_labels",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("name", sa.String, nullable=False, unique=True),
+        sa.Column("description", sa.String, nullable=True),
+        sa.Column("color", sa.String, server_default=sa.text("'#ffffff'"), nullable=False),
+        sa.Column(
+            "user_id",
+            sa.Integer,
+            sa.ForeignKey("users.id", ondelete="SET NULL"),
+            nullable=True,
+        ),
+    )
+
+    op.create_table(
+        "datasets_dataset_labels",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column(
+            "dataset_id",
+            sa.Integer,
+            sa.ForeignKey("datasets.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+        sa.Column(
+            "dataset_label_id",
+            sa.Integer,
+            sa.ForeignKey("dataset_labels.id", ondelete="CASCADE"),
+            nullable=False,
+            index=True,
+        ),
+        sa.UniqueConstraint("dataset_id", "dataset_label_id"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("datasets_dataset_labels")
+    op.drop_table("dataset_labels")

src/phoenix/db/models.py

@@ -1061,6 +1061,9 @@ class Dataset(Base):
     experiment_tags: Mapped[list["ExperimentTag"]] = relationship(
         "ExperimentTag", back_populates="dataset"
     )
+    datasets_dataset_labels: Mapped[list["DatasetsDatasetLabel"]] = relationship(
+        "DatasetsDatasetLabel", back_populates="dataset"
+    )
 
     @hybrid_property
     def example_count(self) -> Optional[int]:
@@ -1111,6 +1114,38 @@ async def load_example_count(self, session: AsyncSession) -> None:
             )
 
 
+class DatasetLabel(Base):
+    __tablename__ = "dataset_labels"
+    name: Mapped[str] = mapped_column(unique=True)
+    description: Mapped[Optional[str]]
+    color: Mapped[str] = mapped_column(String, server_default=text("'#ffffff'"), nullable=False)
+    datasets_dataset_labels: Mapped[list["DatasetsDatasetLabel"]] = relationship(
+        "DatasetsDatasetLabel", back_populates="dataset_label"
+    )
+    user_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("users.id", ondelete="SET NULL"),
+        nullable=True,
+    )
+    user: Mapped[Optional["User"]] = relationship("User")
+
+
+class DatasetsDatasetLabel(Base):
+    __tablename__ = "datasets_dataset_labels"
+    dataset_id: Mapped[int] = mapped_column(
+        ForeignKey("datasets.id", ondelete="CASCADE"),
+    )
+    dataset_label_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_labels.id", ondelete="CASCADE"),
+        index=True,
+    )
+    dataset: Mapped["Dataset"] = relationship("Dataset", back_populates="datasets_dataset_labels")
+    dataset_label: Mapped["DatasetLabel"] = relationship(
+        "DatasetLabel", back_populates="datasets_dataset_labels"
+    )
+
+    __table_args__ = (UniqueConstraint("dataset_id", "dataset_label_id"),)
+
+
 class DatasetVersion(Base):
     __tablename__ = "dataset_versions"
     dataset_id: Mapped[int] = mapped_column(