Setup tag orm

src/database/schema/setups.py

@@ -0,0 +1,16 @@
+"""Setup ORM Model."""
+
+from sqlalchemy.orm import Mapped, mapped_column
+
+from database.schema.base import Base, ExpDBReflected
+from routers.types import Identifier
+
+
+class Setup(ExpDBReflected, Base):
+    """Specifies the hyperparameter configuration of a Flow used in a Run."""
+
+    __tablename__ = "algorithm_setup"
+    id: Mapped[Identifier] = mapped_column("sid", primary_key=True)
+    flow_id: Mapped[Identifier] = mapped_column("implementation_id")
+
+    setup_string: Mapped[str]

src/database/schema/tags.py

@@ -28,3 +28,14 @@ class TaskTag(ExpDBReflected, Tag, Base):
     def task_id(self) -> Identifier:
         """Identifier of the task which is tagged by this tag."""
         return self.entity_id
+
+
+class SetupTag(ExpDBReflected, Tag, Base):
+    """Tags belonging to a setup."""
+
+    __tablename__ = "setup_tag"
+
+    @property
+    def setup_id(self) -> Identifier:
+        """Identifier of the setup which is tagged by this tag."""
+        return self.entity_id

src/database/setups.py

@@ -1,8 +1,9 @@
 """All database operations that directly operate on setups."""
 
+from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from sqlalchemy import text
+from sqlalchemy import select, text
 from sqlalchemy.exc import IntegrityError
 
 from database.exceptions import (
@@ -11,27 +12,18 @@
     DuplicatePrimaryKeyError,
     ForeignKeyConstraintError,
 )
-from database.schema.base import UntypedRow
+from database.schema.setups import Setup
+from database.schema.tags import SetupTag
 from routers.types import Identifier, TagString
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import RowMapping
-    from sqlalchemy.ext.asyncio import AsyncConnection
+    from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 
 
-async def get(setup_id: Identifier, connection: AsyncConnection) -> UntypedRow | None:
+async def get(setup_id: Identifier, session: AsyncSession) -> Setup | None:
     """Get the setup with id `setup_id` from the database."""
-    row = await connection.execute(
-        text(
-            """
-            SELECT *
-            FROM algorithm_setup
-            WHERE sid = :setup_id
-            """,
-        ),
-        parameters={"setup_id": setup_id},
-    )
-    return row.first()
+    return await session.get(Setup, setup_id)
 
 
 async def get_parameters(setup_id: Identifier, connection: AsyncConnection) -> list[RowMapping]:
@@ -61,51 +53,33 @@ async def get_parameters(setup_id: Identifier, connection: AsyncConnection) -> l
     return list(rows.mappings().all())
 
 
-async def get_tags(setup_id: Identifier, connection: AsyncConnection) -> list[UntypedRow]:
+async def get_tags(setup_id: Identifier, session: AsyncSession) -> Sequence[SetupTag]:
     """Get all tags for setup with `setup_id` from the database."""
-    rows = await connection.execute(
-        text(
-            """
-            SELECT *
-            FROM setup_tag
-            WHERE id = :setup_id
-            """,
-        ),
-        parameters={"setup_id": setup_id},
-    )
-    return list(rows.all())
+    stmt = select(SetupTag).where(SetupTag.entity_id == setup_id)
+    return (await session.scalars(stmt)).all()
 
 
-async def untag(setup_id: Identifier, tag: TagString, connection: AsyncConnection) -> None:
-    """Remove tag `tag` from setup with id `setup_id`."""
-    await connection.execute(
-        text(
-            """
-            DELETE FROM setup_tag
-            WHERE id = :setup_id AND tag = :tag
-            """,
-        ),
-        parameters={"setup_id": setup_id, "tag": tag},
-    )
+async def get_tag(setup_id: Identifier, tag: TagString, session: AsyncSession) -> SetupTag | None:
+    """Get the tag `tag` for setup with id `setup_id`."""
+    return await session.get(SetupTag, {"tag": tag, "entity_id": setup_id})
+
+
+async def delete_tag(tag: SetupTag, session: AsyncSession) -> None:
+    """Delete a setup tag."""
+    await session.delete(tag)
 
 
 async def tag(
     setup_id: Identifier,
     tag: TagString,
     user_id: Identifier,
-    connection: AsyncConnection,
+    session: AsyncSession,
 ) -> None:
     """Add tag `tag` to setup with id `setup_id`."""
+    tag_ = SetupTag(entity_id=setup_id, tag=tag, uploader_id=user_id)
     try:
-        await connection.execute(
-            text(
-                """
-                INSERT INTO setup_tag (id, tag, uploader)
-                VALUES (:setup_id, :tag, :user_id)
-                """,
-            ),
-            parameters={"setup_id": setup_id, "tag": tag, "user_id": user_id},
-        )
+        session.add(tag_)
+        await session.flush()
     except IntegrityError as e:
         if e.orig is None:
             raise

src/database/tasks.py

@@ -161,8 +161,8 @@ async def tag(
     user_id: Identifier,
     session: AsyncSession,
 ) -> None:
+    tag = TaskTag(entity_id=id_, uploader_id=user_id, tag=tag_)
     try:
-        tag = TaskTag(entity_id=id_, uploader_id=user_id, tag=tag_)
         session.add(tag)
         await session.flush()
     except IntegrityError as e:

src/routers/dependencies.py

@@ -29,7 +29,7 @@ async def userdb_connection() -> AsyncIterator[AsyncConnection]:
 async def expdb_session(
     connection: Annotated[AsyncConnection, Depends(expdb_connection)],
 ) -> AsyncIterator[AsyncSession]:
-    async with AsyncSession(connection) as session, session.begin():
+    async with AsyncSession(connection) as session:
         yield session
 
 

src/routers/openml/runs.py

@@ -14,7 +14,8 @@
 import database.users
 from core.errors import RunNotFoundError, RunTraceNotFoundError
 from database.schema.base import UntypedRow
-from routers.dependencies import expdb_connection, userdb_connection
+from database.schema.setups import Setup
+from routers.dependencies import expdb_connection, expdb_session, userdb_connection
 from routers.types import Identifier
 from schemas.runs import (
     EvaluationScore,
@@ -28,7 +29,7 @@
 )
 
 if TYPE_CHECKING:
-    from sqlalchemy.ext.asyncio import AsyncConnection
+    from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 
 router = APIRouter(prefix="/run", tags=["run"])
 
@@ -75,15 +76,15 @@ class RunContext:
     evaluation_rows: list[UntypedRow]
     task_type: str | None
     task_evaluation_measure: str | None
-    setup: UntypedRow | None
+    setup: Setup | None
     parameter_rows: list[UntypedRow]
 
 
 async def _load_run_context(
     run: UntypedRow,
-    run_id: int,
     expdb: AsyncConnection,
     userdb: AsyncConnection,
+    expdb_session: AsyncSession,
     engine_ids: list[int],
 ) -> RunContext:
     (
@@ -97,16 +98,16 @@ async def _load_run_context(
         setup,
         parameter_rows,
     ) = cast(
-        "tuple[Any, list[str], list[UntypedRow], list[UntypedRow], list[UntypedRow], str | None, str | None, UntypedRow | None, list[UntypedRow]]",  # noqa: E501
+        "tuple[Any, list[str], list[UntypedRow], list[UntypedRow], list[UntypedRow], str | None, str | None, Setup | None, list[UntypedRow]]",  # noqa: E501
         await asyncio.gather(
             database.users.get_user(user_id=run.uploader, connection=userdb),
-            database.runs.get_tags(run_id, expdb),
-            database.runs.get_input_data(run_id, expdb),
-            database.runs.get_output_files(run_id, expdb),
-            database.runs.get_evaluations(run_id, expdb, evaluation_engine_ids=engine_ids),
+            database.runs.get_tags(run.rid, expdb),
+            database.runs.get_input_data(run.rid, expdb),
+            database.runs.get_output_files(run.rid, expdb),
+            database.runs.get_evaluations(run.rid, expdb, evaluation_engine_ids=engine_ids),
             database.tasks.get_task_type_name(run.task_id, expdb),
             database.tasks.get_task_evaluation_measure(run.task_id, expdb),
-            database.setups.get(run.setup, expdb),
+            database.setups.get(run.setup, expdb_session),
             database.setups.get_parameters(run.setup, expdb),
         ),
     )
@@ -146,6 +147,7 @@ async def get_run(
     run_id: int,
     expdb: Annotated[AsyncConnection, Depends(expdb_connection)],
     userdb: Annotated[AsyncConnection, Depends(userdb_connection)],
+    expdb_session: Annotated[AsyncSession, Depends(expdb_session)],
 ) -> Run:
     """Get full metadata for a run by ID.
 
@@ -158,9 +160,9 @@ async def get_run(
         raise RunNotFoundError(msg, code=236)
 
     engine_ids: list[int] = config.get_config().development.run_evaluation_engine_ids
-    ctx = await _load_run_context(run, run_id, expdb, userdb, engine_ids)
+    ctx = await _load_run_context(run, expdb, userdb, expdb_session, engine_ids)
 
-    flow = await database.flows.get(ctx.setup.implementation_id, expdb) if ctx.setup else None
+    flow = await database.flows.get(ctx.setup.flow_id, expdb) if ctx.setup else None
     evaluations = _build_evaluations(ctx.evaluation_rows)
 
     normalised_measure = ctx.task_evaluation_measure or None
@@ -173,7 +175,7 @@ async def get_run(
         task_id=run.task_id,
         task_type=ctx.task_type,
         task_evaluation_measure=normalised_measure,
-        flow_id=ctx.setup.implementation_id if ctx.setup else None,
+        flow_id=ctx.setup.flow_id if ctx.setup else None,
         flow_name=flow.full_name if flow else None,
         setup_id=run.setup,
         setup_string=ctx.setup.setup_string if ctx.setup else None,

src/routers/openml/setups.py

@@ -1,6 +1,5 @@
 """All endpoints that relate to setups."""
 
-import asyncio
 from typing import TYPE_CHECKING, Annotated
 
 from fastapi import APIRouter, Body, Depends, Path
@@ -15,12 +14,12 @@
 )
 from database.exceptions import DuplicatePrimaryKeyError, ForeignKeyConstraintError
 from database.users import User
-from routers.dependencies import expdb_connection, fetch_user_or_raise
+from routers.dependencies import expdb_connection, expdb_session, fetch_user_or_raise
 from routers.types import Identifier, TagString
 from schemas.setups import SetupParameters, SetupResponse
 
 if TYPE_CHECKING:
-    from sqlalchemy.ext.asyncio import AsyncConnection
+    from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
 
 router = APIRouter(prefix="/setup", tags=["setup"])
 
@@ -29,9 +28,10 @@
 async def get_setup(
     setup_id: Annotated[Identifier, Path()],
     expdb_db: Annotated[AsyncConnection, Depends(expdb_connection)],
+    expdb_session: Annotated[AsyncSession, Depends(expdb_session)],
 ) -> SetupResponse:
     """Get setup by id."""
-    setup = await database.setups.get(setup_id, expdb_db)
+    setup = await database.setups.get(setup_id, expdb_session)
     if not setup:
         msg = f"Setup {setup_id} not found."
         raise SetupNotFoundError(msg, code=281)
@@ -40,7 +40,7 @@ async def get_setup(
 
     params_model = SetupParameters(
         setup_id=setup_id,
-        flow_id=setup.implementation_id,
+        flow_id=setup.flow_id,
         parameter=setup_parameters or None,
     )
 
@@ -52,11 +52,11 @@ async def tag_setup(
     setup_id: Annotated[Identifier, Body()],
     tag: Annotated[TagString, Body()],
     user: Annotated[User, Depends(fetch_user_or_raise)],
-    expdb_db: Annotated[AsyncConnection, Depends(expdb_connection)],
+    expdb_session: Annotated[AsyncSession, Depends(expdb_session)],
 ) -> dict[str, dict[str, str | list[str]]]:
     """Add tag `tag` to setup with id `setup_id`."""
     try:
-        await database.setups.tag(setup_id, tag, user.user_id, expdb_db)
+        await database.setups.tag(setup_id, tag, user.user_id, expdb_session)
     except ForeignKeyConstraintError:
         msg = f"Setup {setup_id} not found."
         raise SetupNotFoundError(msg, code=472) from None
@@ -65,7 +65,7 @@ async def tag_setup(
         raise TagAlreadyExistsError(msg) from None
 
     logger.info("Setup {setup_id} tagged '{tag}'.", setup_id=setup_id, tag=tag)
-    all_tag_rows = await database.setups.get_tags(setup_id, expdb_db)
+    all_tag_rows = await database.setups.get_tags(setup_id, expdb_session)
     all_tags = [t.tag for t in all_tag_rows]
 
     return {"setup_tag": {"id": str(setup_id), "tag": all_tags}}
@@ -76,23 +76,21 @@ async def untag_setup(
     setup_id: Annotated[Identifier, Body()],
     tag: Annotated[TagString, Body()],
     user: Annotated[User, Depends(fetch_user_or_raise)],
-    expdb_db: Annotated[AsyncConnection, Depends(expdb_connection)],
+    expdb_session: Annotated[AsyncSession, Depends(expdb_session)],
 ) -> dict[str, dict[str, str | list[str]]]:
     """Remove tag `tag` from setup with id `setup_id`."""
-    setup, setup_tags = await asyncio.gather(
-        database.setups.get(setup_id, expdb_db),
-        database.setups.get_tags(setup_id, expdb_db),
-    )
-    if not setup:
-        msg = f"Setup {setup_id} not found."
-        raise SetupNotFoundError(msg)
-    matched_tag_row = next((t for t in setup_tags if t.tag.casefold() == tag.casefold()), None)
-
-    if not matched_tag_row:
+    # Setups don't really have an owner, they are associated with runs.
+    # So only the tagger or admins can remove the tag.
+    tag_orm = await database.setups.get_tag(setup_id, tag, expdb_session)
+    if not tag_orm:
+        setup = await database.setups.get(setup_id, expdb_session)
+        if not setup:
+            msg = f"Setup {setup_id} not found."
+            raise SetupNotFoundError(msg)
         msg = f"Setup {setup_id} does not have tag {tag!r}."
         raise TagNotFoundError(msg)
 
-    if matched_tag_row.uploader != user.user_id and not await user.is_admin():
+    if tag_orm.uploader_id != user.user_id and not await user.is_admin():
         msg = (
             f"You may not remove tag {tag!r} of setup {setup_id} because it was not created by you."
         )
@@ -103,9 +101,7 @@ async def untag_setup(
         )
         raise TagNotOwnedError(msg)
 
-    await database.setups.untag(setup_id, matched_tag_row.tag, expdb_db)
+    await database.setups.delete_tag(tag_orm, expdb_session)
     logger.info("Setup {setup_id} had tag '{tag}' removed.", setup_id=setup_id, tag=tag)
-    remaining_tags = [
-        t.tag for t in setup_tags if t.tag.casefold() != matched_tag_row.tag.casefold()
-    ]
+    remaining_tags = [t.tag for t in await database.setups.get_tags(setup_id, expdb_session)]
     return {"setup_untag": {"id": str(setup_id), "tag": remaining_tags}}