feat: support for upscaling endpoints (#11)

* docs: updated examples

* docs: updated tiles example

* docs: fixed wording on unit_jobs schema

* feat: setup of upscale task endpoints

* feat: added saving of upscaling tasks to db

* chore: import cleanup

* chore: cleanup imports

* feat: integrated status retrieval f or upscaling tasks

* feat: integration of status retrieval

* feat: added more details to the job summary to prevent sending too many requests

* test: fixed tests

* fix: linting issues

Author: JanssenBrm

alembic/env.py

@@ -7,7 +7,10 @@
 from alembic import context
 
 from app.database.db import Base  # import your Base here
-from app.database.models.processing_job import ProcessingJobRecord  # import your models here
+from app.database.models.processing_job import (
+    ProcessingJobRecord,
+)
+from app.database.models.upscaling_task import UpscalingTaskRecord
 
 
 # this is the Alembic Config object, which provides
@@ -71,13 +74,11 @@ def run_migrations_online() -> None:
         config.get_section(config.config_ini_section, {}),
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
-        url=get_url()
+        url=get_url(),
     )
 
     with connectable.connect() as connection:
-        context.configure(
-            connection=connection, target_metadata=target_metadata
-        )
+        context.configure(connection=connection, target_metadata=target_metadata)
 
         with context.begin_transaction():
             context.run_migrations()

alembic/versions/046ada00e750_creation_of_upscaling_tables.py

@@ -0,0 +1,107 @@
+"""creation of upscaling tables
+
+Revision ID: 046ada00e750
+Revises: 6105be2a2b21
+Create Date: 2025-08-22 13:22:24.266047
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from sqlalchemy.dialects.postgresql import ENUM as PG_ENUM
+
+
+# revision identifiers, used by Alembic.
+revision: str = "046ada00e750"
+down_revision: Union[str, Sequence[str], None] = "6105be2a2b21"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+process_type_enum = PG_ENUM(
+    "OPENEO", "OGC_API_PROCESS", name="processtypeenum", create_type=False
+)
+
+processing_status_enum = PG_ENUM(
+    "CREATED",
+    "QUEUED",
+    "RUNNING",
+    "FINISHED",
+    "CANCELED",
+    "FAILED",
+    "UNKNOWN",
+    name="processingstatusenum",
+    create_type=False,
+)
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "upscaling_tasks",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("title", sa.String(), nullable=False),
+        sa.Column("label", process_type_enum, nullable=False),
+        sa.Column("status", processing_status_enum, nullable=False),
+        sa.Column("user_id", sa.String(), nullable=False),
+        sa.Column("service", sa.String(), nullable=False),
+        sa.Column("created", sa.DateTime(), nullable=False),
+        sa.Column("updated", sa.DateTime(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_created"), "upscaling_tasks", ["created"], unique=False
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_id"), "upscaling_tasks", ["id"], unique=False
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_label"), "upscaling_tasks", ["label"], unique=False
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_service"), "upscaling_tasks", ["service"], unique=False
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_status"), "upscaling_tasks", ["status"], unique=False
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_title"), "upscaling_tasks", ["title"], unique=False
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_updated"), "upscaling_tasks", ["updated"], unique=False
+    )
+    op.create_index(
+        op.f("ix_upscaling_tasks_user_id"), "upscaling_tasks", ["user_id"], unique=False
+    )
+    op.add_column(
+        "processing_jobs", sa.Column("upscaling_task_id", sa.Integer(), nullable=True)
+    )
+    op.create_foreign_key(
+        None,
+        "processing_jobs",
+        "upscaling_tasks",
+        ["upscaling_task_id"],
+        ["id"],
+        ondelete="SET NULL",
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, "processing_jobs", type_="foreignkey")
+    op.drop_column("processing_jobs", "upscaling_task_id")
+    op.drop_index(op.f("ix_upscaling_tasks_user_id"), table_name="upscaling_tasks")
+    op.drop_index(op.f("ix_upscaling_tasks_updated"), table_name="upscaling_tasks")
+    op.drop_index(op.f("ix_upscaling_tasks_title"), table_name="upscaling_tasks")
+    op.drop_index(op.f("ix_upscaling_tasks_status"), table_name="upscaling_tasks")
+    op.drop_index(op.f("ix_upscaling_tasks_service"), table_name="upscaling_tasks")
+    op.drop_index(op.f("ix_upscaling_tasks_label"), table_name="upscaling_tasks")
+    op.drop_index(op.f("ix_upscaling_tasks_id"), table_name="upscaling_tasks")
+    op.drop_index(op.f("ix_upscaling_tasks_created"), table_name="upscaling_tasks")
+    op.drop_table("upscaling_tasks")
+    # ### end Alembic commands ###

app/database/models/processing_job.py

@@ -1,11 +1,12 @@
 import datetime
 from typing import List, Optional
+
 from loguru import logger
-from sqlalchemy import DateTime, Enum, Integer, String
-from app.database.db import Base
-from sqlalchemy.orm import Session, Mapped, mapped_column
+from sqlalchemy import DateTime, Enum, ForeignKey, Integer, String
+from sqlalchemy.orm import Mapped, Session, mapped_column
 
-from app.schemas.unit_job import ProcessTypeEnum, ProcessingStatusEnum
+from app.database.db import Base
+from app.schemas.unit_job import ProcessingStatusEnum, ProcessTypeEnum
 
 
 class ProcessingJobRecord(Base):
@@ -34,6 +35,12 @@ class ProcessingJobRecord(Base):
         index=True,
     )
 
+    upscaling_task_id: Mapped[int | None] = mapped_column(
+        Integer,
+        ForeignKey("upscaling_tasks.id", ondelete="SET NULL"),
+        nullable=True,
+    )
+
 
 def save_job_to_db(
     db_session: Session, job: ProcessingJobRecord
@@ -51,11 +58,22 @@ def save_job_to_db(
     return job
 
 
-def get_jobs_by_user_id(database: Session, user_id: str) -> List[ProcessingJobRecord]:
-    logger.info(f"Retrieving all processing jobs for user {user_id}")
+def get_jobs_by_user_id(
+    database: Session, user_id: str, upscaling_task_id: Optional[int]
+) -> List[ProcessingJobRecord]:
+    logger.info(
+        f"Retrieving all processing jobs for user {user_id} for upscaling task {upscaling_task_id}"
+    )
     return (
         database.query(ProcessingJobRecord)
-        .filter(ProcessingJobRecord.user_id == user_id)
+        .filter(
+            ProcessingJobRecord.user_id == user_id,
+            (
+                ProcessingJobRecord.upscaling_task_id == upscaling_task_id
+                if upscaling_task_id
+                else ProcessingJobRecord.upscaling_task_id.is_(None)
+            ),
+        )
         .all()
     )
 

app/database/models/upscaling_task.py

@@ -0,0 +1,113 @@
+import datetime
+from typing import List, Optional
+
+from loguru import logger
+from sqlalchemy import DateTime, Enum, Integer, String
+from sqlalchemy.orm import Mapped, Session, mapped_column
+
+from app.database.db import Base
+from app.schemas.unit_job import ProcessingStatusEnum, ProcessTypeEnum
+
+
+class UpscalingTaskRecord(Base):
+    __tablename__ = "upscaling_tasks"
+
+    id: Mapped[int] = mapped_column(
+        Integer, primary_key=True, index=True, autoincrement=True
+    )
+    title: Mapped[str] = mapped_column(String, index=True)
+    label: Mapped[ProcessTypeEnum] = mapped_column(Enum(ProcessTypeEnum), index=True)
+    status: Mapped[ProcessingStatusEnum] = mapped_column(
+        Enum(ProcessingStatusEnum), index=True
+    )
+    user_id: Mapped[str] = mapped_column(String, index=True)
+    service: Mapped[str] = mapped_column(String, index=True)
+    created: Mapped[datetime.datetime] = mapped_column(
+        DateTime, default=datetime.datetime.utcnow, index=True
+    )
+    updated: Mapped[datetime.datetime] = mapped_column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        index=True,
+    )
+
+
+def save_upscaling_task_to_db(
+    db_session: Session, task: UpscalingTaskRecord
+) -> UpscalingTaskRecord:
+    """
+    Save an upscaling task record to the database and update the ID of the task.
+
+    :param db_session: The database session to use for saving the task.
+    :param job: The UpscalingTaskRecord instance to save.
+    """
+    db_session.add(task)
+    db_session.commit()
+    db_session.refresh(task)  # Refresh to get the ID after commit
+    logger.debug(f"Upscale task saved with ID: {task.id}")
+    return task
+
+
+def get_upscale_tasks_by_user_id(
+    database: Session, user_id: str
+) -> List[UpscalingTaskRecord]:
+    logger.info(f"Retrieving all upscale tasks for user {user_id}")
+    return (
+        database.query(UpscalingTaskRecord)
+        .filter(UpscalingTaskRecord.user_id == user_id)
+        .all()
+    )
+
+
+def get_upscale_task_by_id(
+    database: Session, task_id: int
+) -> Optional[UpscalingTaskRecord]:
+    logger.info(f"Retrieving upscale task with ID {task_id}")
+    return (
+        database.query(UpscalingTaskRecord)
+        .filter(UpscalingTaskRecord.id == task_id)
+        .first()
+    )
+
+
+def get_upscale_task_by_user_id(
+    database: Session, task_id: int, user_id: str
+) -> Optional[UpscalingTaskRecord]:
+    logger.info(f"Retrieving upscale task with ID {task_id} for user {user_id}")
+    return (
+        database.query(UpscalingTaskRecord)
+        .filter(
+            UpscalingTaskRecord.id == task_id, UpscalingTaskRecord.user_id == user_id
+        )
+        .first()
+    )
+
+
+def update_upscale_task_status_by_id(
+    database: Session, task_id: int, status: ProcessingStatusEnum
+):
+    logger.info(f"Updating the status of upscale task with ID {task_id} to {status}")
+    task = get_upscale_task_by_id(database, task_id)
+
+    if task:
+        task.status = status
+        database.commit()
+        database.refresh(task)
+    else:
+        logger.warning(
+            f"Could not update upscaling task status of task {task_id} as it could not be found "
+            "in the database"
+        )
+
+
+# def get_jobs_by_user_id(database: Session, user_id: str) -> List[ProcessingJobRecord]:
+#     logger.info(f"Retrieving all processing jobs for user {user_id}")
+#     return (
+#         database.query(ProcessingJobRecord)
+#         .filter(
+#             ProcessingJobRecord.user_id == user_id,
+#             ProcessingJobRecord.upscaling_task_id is None,
+#         )
+#         .all()
+#     )

app/main.py

@@ -5,7 +5,7 @@
 from app.services.tiles.base import load_grids
 from .config.logger import setup_logging
 from .config.settings import settings
-from .routers import jobs_status, unit_jobs, health, tiles
+from .routers import jobs_status, unit_jobs, health, tiles, upscale_tasks
 
 setup_logging()
 
@@ -27,4 +27,5 @@
 app.include_router(tiles.router)
 app.include_router(jobs_status.router)
 app.include_router(unit_jobs.router)
+app.include_router(upscale_tasks.router)
 app.include_router(health.router)

app/routers/jobs_status.py

@@ -7,6 +7,7 @@
 from app.database.db import get_db
 from app.schemas.jobs_status import JobsStatusResponse
 from app.services.processing import get_processing_jobs_by_user_id
+from app.services.upscaling import get_upscaling_tasks_by_user_id
 
 router = APIRouter()
 
@@ -25,7 +26,8 @@ async def get_jobs_status(
     """
     logger.debug(f"Fetching jobs list for user {user}")
     return JobsStatusResponse(
-        upscaling_tasks=[], processing_jobs=get_processing_jobs_by_user_id(db, user)
+        upscaling_tasks=get_upscaling_tasks_by_user_id(db, user),
+        processing_jobs=get_processing_jobs_by_user_id(db, user),
     )