dodil/erp-platform

tests/scenario_bricks/scn_orders/infrastructure/tables.py

"""orders' tables: an order (versioned) with its lines, and the manual tasks its steps leave for people."""
from typing import Optional

from sqlalchemy import Boolean, Column, Integer, MetaData, Numeric, String, Table, TIMESTAMP, select

from erp_platform.repository import Aggregate, Part

md = MetaData()
orders = Table("orders_orders", md,
               Column("order_id", String, primary_key=True), Column("customer", String), Column("owner", String),
               Column("status", String), Column("currency", String), Column("placed_at", TIMESTAMP),
               Column("due_at", TIMESTAMP), Column("reserved", Boolean), Column("reserve_note", String),
               Column("note", String), Column("created_at", TIMESTAMP), Column("updated_at", TIMESTAMP),
               Column("version", Integer))
lines = Table("orders_lines", md,
              Column("order_id", String, primary_key=True), Column("line_no", Integer, primary_key=True),
              Column("sku", String), Column("quantity", Integer), Column("unit_price", Numeric(18, 2)),
              Column("created_at", TIMESTAMP), Column("updated_at", TIMESTAMP))
tasks = Table("orders_tasks", md,
              Column("task_key", String, primary_key=True), Column("order_id", String), Column("capability", String),
              Column("input_json", String), Column("plan", String), Column("reason", String), Column("state", String),
              Column("created_at", TIMESTAMP), Column("updated_at", TIMESTAMP))

ORDER = Aggregate(Part(orders, ("order_id",)), Part(lines, ("order_id", "line_no")), version="version")
TASK = Aggregate(Part(tasks, ("task_key",)))


class OrderRepository:
    def __init__(self, conn):
        self.conn = conn

    def get(self, order_id: str) -> Optional[dict]:
        rows = ORDER.load(self.conn, {"order_id": order_id})
        if rows is None:
            return None
        head = dict(rows["orders_orders"][0])
        head["lines"] = sorted(rows["orders_lines"], key=lambda r: r["line_no"])
        return head

    def save(self, order: dict, expected_version: int) -> int:
        head = {k: v for k, v in order.items() if k in orders.c and k not in ("created_at", "updated_at")}
        return ORDER.save(self.conn, {
            "orders_orders": [head],
            "orders_lines": [{k: v for k, v in ln.items() if k in lines.c and k not in ("created_at", "updated_at")}
                             for ln in order["lines"]]}, expected_version=expected_version)

    def ids(self, owner: Optional[str]) -> list:
        stmt = select(orders.c.order_id)
        if owner is not None:
            stmt = stmt.where(orders.c.owner == owner)
        return sorted(r[0] for r in self.conn.execute(stmt))

    def open_due(self, now) -> list:
        return sorted(r[0] for r in self.conn.execute(
            select(orders.c.order_id).where(orders.c.status == "open", orders.c.due_at <= now.replace(tzinfo=None))))

    def overdue(self) -> list:
        return sorted(r[0] for r in self.conn.execute(select(orders.c.order_id).where(orders.c.status == "overdue")))

    def task(self, key: str) -> Optional[dict]:
        rows = TASK.load(self.conn, {"task_key": key})
        return dict(rows["orders_tasks"][0]) if rows else None

    def add_task(self, row: dict) -> None:
        TASK.save(self.conn, {"orders_tasks": [row]})

    def tasks_for(self, order_ids: list) -> list:
        if not order_ids:
            return []
        return [dict(r._mapping) for r in self.conn.execute(
            select(tasks).where(tasks.c.order_id.in_(order_ids)).order_by(tasks.c.task_key))]