Model/backend/app/documents/download_tasks.py

"""The bulk-download route's view of the app-owned task's sub_task (ADR-0060).

Parameterised SQL rather than the SQLModel mirrors: importing the
``infrastructure.postgres`` mirror of ``sub_task`` into the FastAPI process
double-registers the table and crashes at import (the same constraint the
Modelling Run distributor works around — see ``modelling/run_tasks.py``).
"""

import json
from datetime import datetime, timezone
from typing import Any, cast
from uuid import UUID

from sqlalchemy import text
from sqlmodel import Session


class DocumentDownloadTasks:
    def __init__(self, session: Session) -> None:
        self._session = session

    def read_selection_config(self, task_id: UUID) -> dict[str, Any]:
        """The task's selection config, read from the FE-owned ``task.inputs``
        JSON (ADR-0060) — ``{portfolio_id, property_ids?, select_all?}``. Empty
        dict when the task has no inputs."""
        row = (
            self._session.connection()
            .execute(
                text("SELECT inputs FROM tasks WHERE id = :task_id"),
                {"task_id": task_id},
            )
            .first()
        )
        if row is None or row[0] is None:
            return {}
        parsed: Any = json.loads(row[0])
        return cast(dict[str, Any], parsed) if isinstance(parsed, dict) else {}

    def already_distributed(self, task_id: UUID) -> bool:
        """Whether the task already has a sub_task — the 409 guard against a
        double-submit re-triggering the same download."""
        row = (
            self._session.connection()
            .execute(
                text("SELECT 1 FROM sub_task WHERE task_id = :task_id LIMIT 1"),
                {"task_id": task_id},
            )
            .first()
        )
        return row is not None

    def resolve_selection(
        self,
        project_codes: list[str],
        property_ids: list[int],
    ) -> list[str]:
        """Resolve the selection to the distinct ``landlord_property_id`` set the
        Download Package is built from (ADR-0060). Two independent sources,
        unioned:

        - **project codes** — every property in the named HubSpot projects, read
          from ``hubspot_deal_data`` (the project↔property grain lives there, not
          on ``property``).
        - **hand-picked property ids** — the selected ``property`` rows.

        Files are matched by ``landlord_property_id`` (the missing ``property_id``
        on ``uploaded_files`` is a known future gap), so a source row with a null
        ``landlord_property_id`` contributes nothing and is dropped here. The
        returned set size is the cap basis (property count, ADR-0060).
        """
        landlord_property_ids: set[str] = set()
        if project_codes:
            rows = self._session.connection().execute(
                text(
                    "SELECT DISTINCT landlord_property_id FROM hubspot_deal_data"
                    " WHERE project_code = ANY(:codes)"
                    " AND landlord_property_id IS NOT NULL"
                ),
                {"codes": project_codes},
            )
            landlord_property_ids.update(row[0] for row in rows.all())
        if property_ids:
            rows = self._session.connection().execute(
                text(
                    "SELECT landlord_property_id FROM property"
                    " WHERE id = ANY(:ids) AND landlord_property_id IS NOT NULL"
                ),
                {"ids": property_ids},
            )
            landlord_property_ids.update(row[0] for row in rows.all())
        return sorted(landlord_property_ids)

    def create_download_subtask(
        self, task_id: UUID, subtask_id: UUID, recipe: dict[str, Any]
    ) -> bool:
        """Pre-create the single ``waiting`` sub_task under the app-owned task,
        pinning the recipe (landlord_property_ids, recipient_email,
        package_name) onto its inputs — the worker's reproducible instructions
        (ADR-0055/0060).

        The insert is **conditional on the task having no sub_task yet**, so the
        DB — not a prior read — arbitrates a double-submit race: two concurrent
        requests can both pass the ``already_distributed`` check, but only one
        insert lands. Returns whether this call created the sub_task."""
        now = datetime.now(timezone.utc)
        result = self._session.connection().execute(
            text(
                "INSERT INTO sub_task (id, task_id, status, inputs, updated_at)"
                " SELECT :id, :task_id, 'waiting', :inputs, :updated_at"
                " WHERE NOT EXISTS ("
                "   SELECT 1 FROM sub_task WHERE task_id = :task_id"
                " )"
            ),
            {
                "id": subtask_id,
                "task_id": task_id,
                "inputs": json.dumps(recipe),
                "updated_at": now,
            },
        )
        self._session.commit()
        return result.rowcount == 1