strata.py

import time
import os

from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route
from starlette.exceptions import HTTPException

from sqlalchemy import create_engine, Column, Integer, String, Float, ForeignKey, select
from sqlalchemy.orm import Session, declarative_base

DATABASE_URL = os.environ.get("DATABASE_URL", "sqlite:///strata.db")
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
Base = declarative_base()


class Period(Base):
    __tablename__ = "periods"

    id = Column(Integer, primary_key=True)
    name = Column(String, nullable=False)
    path = Column(String, unique=True, nullable=False)
    start = Column(Float, nullable=False)
    end = Column(Float, nullable=True)
    parent_id = Column(Integer, ForeignKey("periods.id"), nullable=True)


Base.metadata.create_all(engine)


def get_db():
    return Session(engine)


def _end_active_descendants(db, period, now):
    children = db.execute(
        select(Period).where(Period.parent_id == period.id, Period.end.is_(None))
    ).scalars().all()
    for child in children:
        _end_active_descendants(db, child, now)
        child.end = now


async def handle_get(request):
    path = request.path_params.get("path", "").rstrip("/")

    with get_db() as db:
        if not path:
            roots = db.execute(
                select(Period).where(Period.parent_id.is_(None))
            ).scalars().all()
            return JSONResponse([r.name for r in roots])

        period = db.execute(
            select(Period).where(Period.path == path)
        ).scalar_one_or_none()
        if period is None:
            raise HTTPException(404, "Period not found")

        children = db.execute(
            select(Period).where(Period.parent_id == period.id)
        ).scalars().all()

        return JSONResponse({
            "start": int(period.start),
            "end": int(period.end) if period.end is not None else None,
            "children": [c.name for c in children],
        })


def ensure_parent(db, parent_path, now):
    parent = db.execute(
        select(Period).where(Period.path == parent_path)
    ).scalar_one_or_none()
    if parent is not None:
        if parent.end is not None:
            raise HTTPException(400, f"Ancestor {parent_path} has already ended")
        return parent

    parts = parent_path.split("/")
    name = parts[-1]
    ancestor_path = "/".join(parts[:-1]) if len(parts) > 1 else None

    grandparent = None
    if ancestor_path:
        grandparent = ensure_parent(db, ancestor_path, now)

    if grandparent:
        active_sibling = _find_active_sibling(db, grandparent)
        if active_sibling is not None:
            _end_active_descendants(db, active_sibling, now)
            active_sibling.end = now

    parent = Period(
        name=name,
        path=parent_path,
        start=now,
        parent_id=grandparent.id if grandparent else None,
    )
    db.add(parent)
    return parent


def _find_active_sibling(db, parent_period):
    db.flush()
    return db.execute(
        select(Period).where(
            Period.parent_id == parent_period.id,
            Period.end.is_(None),
        )
    ).scalar_one_or_none()


async def handle_put(request):
    path = request.path_params.get("path", "").rstrip("/")
    if not path:
        raise HTTPException(400, "Path must not be empty")

    parts = path.split("/")
    name = parts[-1]
    parent_path = "/".join(parts[:-1]) if len(parts) > 1 else None

    now = time.time()

    with get_db() as db:
        existing = db.execute(
            select(Period).where(Period.path == path)
        ).scalar_one_or_none()
        if existing is not None:
            if existing.end is not None:
                raise HTTPException(400, "Period has already ended")
            children = db.execute(
                select(Period).where(Period.parent_id == existing.id)
            ).scalars().all()
            return JSONResponse({
                "start": int(existing.start),
                "end": None,
                "children": [c.name for c in children],
            })

        parent = None
        if parent_path:
            parent = ensure_parent(db, parent_path, now)

        if parent:
            active_sibling = _find_active_sibling(db, parent)

            if active_sibling is not None and active_sibling.path != path:
                _end_active_descendants(db, active_sibling, now)
                active_sibling.end = now

        period = Period(
            name=name,
            path=path,
            start=now,
            parent_id=parent.id if parent else None,
        )
        db.add(period)
        db.commit()

        return JSONResponse({
            "start": int(period.start),
            "end": None,
            "children": [],
        }, status_code=201)


async def handle_delete(request):
    path = request.path_params.get("path", "").rstrip("/")
    if not path:
        raise HTTPException(400, "Path must not be empty")

    now = time.time()

    with get_db() as db:
        period = db.execute(
            select(Period).where(Period.path == path)
        ).scalar_one_or_none()
        if period is None:
            raise HTTPException(404, "Period not found")
        if period.end is not None:
            raise HTTPException(400, "Period already ended")

        _end_active_descendants(db, period, now)
        period.end = now
        db.commit()

        return JSONResponse({"status": "ok"})


routes = [
    Route("/{path:path}", endpoint=handle_get, methods=["GET"]),
    Route("/{path:path}", endpoint=handle_put, methods=["PUT"]),
    Route("/{path:path}", endpoint=handle_delete, methods=["DELETE"]),
]

app = Starlette(routes=routes)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)