Accounting endpoints

tests/api/conftest.py

@@ -1,5 +1,6 @@
 import json
 from contextlib import contextmanager
+from datetime import datetime, timedelta
 from typing import AsyncGenerator, Callable, Generator
 
 import pytest
@@ -13,6 +14,7 @@
 from warden.api.routes.dependencies.qpu_client import get_qpu_client
 from warden.lib.config.config import APIConfig, Config, DatabaseConfig, QPUConfig
 from warden.lib.db.database import Base
+from warden.lib.models import Job, Session
 from warden.lib.qpu_client.client import AsyncQPUClient
 
 
@@ -249,3 +251,54 @@ def qpu_specs() -> dict:
     specs = json.loads(DigitalAnalogDevice.to_abstract_repr())
     specs["name"] = "FRESNEL_CAN1"
     return specs
+
+
+async def acct_populate_db(
+    app,
+    serialized_sequence: str,
+    n_users: int,
+    first_session_start: datetime = datetime(2026, 1, 1, 0, 0, 0),
+    session_duration: timedelta = timedelta(hours=1),
+    user_time_offset: timedelta = timedelta(hours=1),
+) -> tuple[list[str], list[Job], list[Session]]:
+    """Creates mock data for accounting testing in DB"""
+    BASE_START_DATETIME = first_session_start
+    BASE_END_DATETIME = BASE_START_DATETIME + session_duration
+    user_uids = [str(i) for i in range(1000, 1000 + n_users)]
+
+    # Create at least one session and job per user
+    sessions = []
+    jobs = []
+    for i, uid in enumerate(user_uids):
+        start_session = BASE_START_DATETIME + i * user_time_offset
+        end_session = BASE_END_DATETIME + i * user_time_offset
+
+        sessions.append(
+            Session(
+                created_at=start_session,
+                revoked_at=end_session,
+                user_id=uid,
+                slurm_job_id=str(i),
+            )
+        )
+        jobs.append(
+            Job(
+                status="DONE",
+                logs="",
+                shots=100,
+                sequence=serialized_sequence,
+                created_at=start_session,
+                scheduled_at=start_session,
+                started_at=start_session,
+                ended_at=end_session,
+                # Relationship
+                session=sessions[-1],
+            )
+        )
+
+    async_session_factory = app.state.db_session_factory
+    async with async_session_factory() as db_session:
+        db_session.add_all(sessions)
+        db_session.add_all(jobs)
+        await db_session.commit()
+    return user_uids, jobs, sessions

tests/api/test_acct.py

@@ -0,0 +1,161 @@
+from datetime import datetime, timedelta
+
+import pytest
+from httpx import AsyncClient
+
+from tests.api.conftest import acct_populate_db, mock_munge_auth
+
+ACCT_ENDPOINT = "/acct"
+
+
+@pytest.mark.asyncio
+async def test_acct_required_start(client: AsyncClient, app):
+    """Assert that 'start_datetime' is required in the accounting data query."""
+
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(ACCT_ENDPOINT)
+    assert response.status_code == 422
+
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(ACCT_ENDPOINT + "?end_datetime=2022-01-01T00:00:00")
+    assert response.status_code == 422
+
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(
+            ACCT_ENDPOINT + "?start_datetime=2022-01-01T00:00:00"
+        )
+    assert response.status_code == 200
+
+
+@pytest.mark.asyncio
+async def test_acct_nominal_pagination_filter(client, app, serialized_sequence):
+    """Assert that the accounting data query returns the right number of rows."""
+    N_USERS = 10
+
+    user_uids, _, _ = await acct_populate_db(app, serialized_sequence, N_USERS)
+
+    # Test default
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(
+            ACCT_ENDPOINT + "?start_datetime=2020-01-01T00:00:00&limit=100"
+        )
+    assert response.status_code == 200
+    assert response.json()["pagination"]["total"] == N_USERS
+    assert response.json()["pagination"]["start"] == 0
+    assert response.json()["pagination"]["end"] == N_USERS
+    assert len(response.json()["data"]) == N_USERS
+    assert response.json()["data"][0]["user_id"] == user_uids[0]
+
+    # Test pagination limit
+    PAGINATION_LIMIT = 5
+    assert PAGINATION_LIMIT < N_USERS
+
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(
+            ACCT_ENDPOINT
+            + f"?start_datetime=2020-01-01T00:00:00&limit={PAGINATION_LIMIT}"
+        )
+    assert response.status_code == 200
+    assert response.json()["pagination"]["total"] == N_USERS
+    assert response.json()["pagination"]["start"] == 0
+    assert response.json()["pagination"]["end"] == PAGINATION_LIMIT
+    assert len(response.json()["data"]) == PAGINATION_LIMIT
+    assert response.json()["data"][0]["user_id"] == user_uids[0]
+
+    # Test offset
+    PAGINATION_LIMIT = 5
+    PAGINATION_OFFSET = 4
+    assert PAGINATION_LIMIT + PAGINATION_OFFSET < N_USERS
+
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(
+            ACCT_ENDPOINT
+            + f"?start_datetime=2020-01-01T00:00:00&limit={PAGINATION_LIMIT}&offset={PAGINATION_OFFSET}"
+        )
+    assert response.status_code == 200
+    assert response.json()["pagination"]["total"] == N_USERS
+    assert response.json()["pagination"]["start"] == PAGINATION_OFFSET
+    assert response.json()["pagination"]["end"] == PAGINATION_LIMIT + PAGINATION_OFFSET
+    assert len(response.json()["data"]) == PAGINATION_LIMIT
+    assert response.json()["data"][0]["user_id"] == user_uids[PAGINATION_OFFSET]
+
+    # Test offset end of data length
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(
+            ACCT_ENDPOINT + f"?start_datetime=2020-01-01T00:00:00&offset={N_USERS - 2}"
+        )
+    assert response.status_code == 200
+    assert response.json()["pagination"]["total"] == N_USERS
+    assert response.json()["pagination"]["start"] == N_USERS - 2
+    assert response.json()["pagination"]["end"] == N_USERS
+    assert len(response.json()["data"]) == 2
+    assert response.json()["data"][0]["user_id"] == user_uids[N_USERS - 2]
+
+
+@pytest.mark.asyncio
+async def test_acct_nominal_datetime_filter(client, app, serialized_sequence):
+    """Assert that the accounting data query correctly filters according to start/end datetime."""
+
+    N_USERS = 10
+    FIRST_SESSION_START = datetime(2026, 1, 1, 0, 0, 0)
+
+    SESSION_DURATION = timedelta(minutes=45)
+    USER_TIME_OFFSET = timedelta(hours=1)
+
+    user_ids, _, sessions = await acct_populate_db(
+        app,
+        serialized_sequence,
+        first_session_start=FIRST_SESSION_START,
+        n_users=N_USERS,
+        session_duration=SESSION_DURATION,
+        user_time_offset=USER_TIME_OFFSET,
+    )
+
+    # Test default get all
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(
+            ACCT_ENDPOINT + "?start_datetime=" + FIRST_SESSION_START.isoformat()
+        )
+    assert response.status_code == 200
+    assert len(response.json()["data"]) == N_USERS
+    assert response.json()["data"][0]["user_id"] == user_ids[0]
+
+    # Test filtering only first user session
+    first_session = sessions[0]
+    assert first_session.revoked_at is not None
+    with mock_munge_auth(app, uid=0):
+        response = await client.get(
+            ACCT_ENDPOINT
+            + "?start_datetime="
+            + FIRST_SESSION_START.isoformat()
+            + "&end_datetime="
+            + (first_session.revoked_at + timedelta(seconds=1)).isoformat()
+        )
+    assert response.status_code == 200
+    assert len(response.json()["data"]) == 1
+    assert response.json()["data"][0]["user_id"] == user_ids[0]
+
+    # Test time filtering is total partition of data
+    # We split the filtering exactly at the revocation time of the 2nd user's session
+    second_session = sessions[1]
+    assert second_session.revoked_at is not None
+    with mock_munge_auth(app, uid=0):
+        response_before = await client.get(
+            ACCT_ENDPOINT
+            + "?start_datetime="
+            + FIRST_SESSION_START.isoformat()
+            + "&end_datetime="
+            + second_session.revoked_at.isoformat()
+            + "&limit=100"
+        )
+        response_after = await client.get(
+            ACCT_ENDPOINT
+            + "?start_datetime="
+            + second_session.revoked_at.isoformat()
+            + "&limit=100"
+        )
+    assert response_after.status_code == 200
+    assert response_before.status_code == 200
+
+    assert len(response_before.json()["data"]) == 1
+    assert len(response_after.json()["data"]) == 9

warden/api/app.py

@@ -2,7 +2,7 @@
 
 from fastapi import FastAPI
 
-from warden.api.routes import accessible, jobs, qpu, sessions
+from warden.api.routes import accessible, acct, jobs, qpu, sessions
 from warden.api.routes.dependencies.auth import init_auth
 from warden.api.routes.dependencies.db import init_db
 from warden.api.routes.dependencies.qpu_client import init_qpu_client
@@ -23,6 +23,7 @@ def create_app(config: Config):
     app.include_router(sessions.router, tags=["sessions"])
     app.include_router(qpu.router, tags=["qpu"])
     app.include_router(accessible.router, tags=["accessible"])
+    app.include_router(acct.router, tags=["accounting"])
 
     logger = logging.getLogger(__name__)
 

warden/api/routes/acct.py

@@ -0,0 +1,141 @@
+"""Accounting information route"""
+
+from logging import getLogger
+from typing import Annotated
+
+from fastapi import APIRouter, Query
+from sqlalchemy import and_, func, select
+
+from warden.api.routes.dependencies.auth import (
+    AdminUserDep,
+)
+from warden.api.routes.dependencies.db import DBSessionDep
+from warden.api.schemas.acct import (
+    AcctData,
+    GetAcctRequest,
+    GetAcctResponse,
+    JobsSummary,
+    JobSummaryStats,
+    PaginationResponse,
+    SessionsSummary,
+)
+from warden.lib.models import Job, Session
+
+logger = getLogger(__name__)
+router = APIRouter(prefix="/acct")
+
+
+@router.get("")
+async def get_accounting_snapshot(
+    params: Annotated[GetAcctRequest, Query()],
+    db_session: DBSessionDep,
+    _admin: AdminUserDep,
+) -> GetAcctResponse:
+    """High-level accounting data"""
+
+    # Base session filter
+    session_filters = [Session.revoked_at >= params.start_datetime]
+
+    if params.end_datetime:
+        session_filters.append(Session.revoked_at < params.end_datetime)
+
+    if params.user_ids:
+        session_filters.append(Session.user_id.in_(params.user_ids))
+
+    # Get total count for pagination
+    total_count_stmt = select(func.count(func.distinct(Session.user_id))).where(
+        and_(*session_filters)
+    )
+    total_result = await db_session.execute(total_count_stmt)
+    total_count = total_result.scalar() or 0
+
+    # Query to get sessions data aggregated by user_id
+    sessions_stmt = (
+        select(
+            Session.user_id,
+            func.count(Session.id).label("session_count"),
+            # Check
+            func.coalesce(
+                func.sum(
+                    func.extract("epoch", Session.revoked_at - Session.created_at)
+                ),
+                0,
+            ).label("total_session_duration"),
+        )
+        .where(and_(*session_filters))
+        .group_by(Session.user_id)
+        .order_by(Session.user_id)
+        .offset(params.offset)
+        .limit(params.limit)
+    )
+
+    sessions_result = await db_session.execute(sessions_stmt)
+    sessions_data = sessions_result.fetchall()
+
+    user_sessions_summary: dict[str, SessionsSummary] = {}
+
+    for session_row in sessions_data:
+        user_id = session_row.user_id
+        user_sessions_summary[user_id] = SessionsSummary(
+            count=session_row.session_count,
+            total_duration=session_row.total_session_duration,
+        )
+
+    # Query to get jobs data aggregated by user_id and job status
+    jobs_stmt = (
+        select(
+            Session.user_id,
+            Job.status,
+            func.count(Job.id).label("job_count"),
+            func.coalesce(
+                func.sum(func.extract("epoch", Job.ended_at - Job.started_at)), 0
+            ).label("total_job_duration"),
+        )
+        .join(Session, Session.id == Job.session_id, isouter=True)
+        .where(and_(*session_filters))
+        .group_by(Session.user_id, Job.status)
+        .order_by(Session.user_id)
+        .offset(params.offset)
+        .limit(params.limit)
+    )
+
+    jobs_result = await db_session.execute(jobs_stmt)
+    jobs_data = jobs_result.fetchall()
+
+    user_jobs_summaries: dict[str, JobsSummary] = {}
+
+    for job_row in jobs_data:
+        user_id = job_row.user_id
+        if user_id not in user_jobs_summaries:
+            user_jobs_summaries[user_id] = JobsSummary(count=0)
+        user_jobs_summaries[user_id].count += job_row.job_count
+        user_jobs_summaries[user_id].stats.append(
+            JobSummaryStats(
+                status=job_row.status,
+                count=job_row.job_count,
+                total_duration=int(job_row.total_job_duration or 0),
+            )
+        )
+
+    # For each user, get job statistics
+    acct_data_list = []
+
+    for user_id in user_sessions_summary.keys():
+        user_acct_data = AcctData(
+            user_id=user_id,
+            sessions=user_sessions_summary[user_id],
+            jobs=user_jobs_summaries.get(user_id, JobsSummary(count=0)),
+        )
+
+        acct_data_list.append(user_acct_data)
+
+    pagination_resp = PaginationResponse(
+        total=total_count,
+        start=params.offset,
+        end=min(params.offset + len(acct_data_list), total_count),
+    )
+
+    return GetAcctResponse(
+        data=acct_data_list,
+        pagination=pagination_resp,
+    )