LedgerLine/services/ledger/repository.py

import asyncio
import uuid

import sqlalchemy as sa

KIND_SIGN = {"deposit": 1, "withdrawal": -1, "purchase": -1}

ROW_LOCK_HOLD_SECONDS = 0.15


async def get_account(session, account_id: str):
    result = await session.execute(
        sa.text("SELECT id, holder_name, currency, balance_minor, is_hot FROM accounts WHERE id = :id"),
        {"id": account_id},
    )
    return result.first()


async def list_accounts(session, limit: int = 1000):
    result = await session.execute(
        sa.text(
            "SELECT id, holder_name, currency, balance_minor FROM accounts ORDER BY created_at LIMIT :limit"
        ),
        {"limit": limit},
    )
    return result.mappings().all()


async def get_transaction_currency(session, tx_id: str):
    """Backs the poison_data pathology: fetches one known-malformed transaction
    (NULL currency, §8) by primary key -- a cheap, always-fast lookup regardless
    of the deliberately-missing account_id index, since which transaction to
    fetch is looked up by account_id in Redis first (chaos.py's
    find_poison_transaction_id), not by scanning `transactions` here."""
    result = await session.execute(
        sa.text("SELECT currency FROM transactions WHERE id = :id"),
        {"id": tx_id},
    )
    return result.first()


async def create_transaction(
    session, *, account_id, kind, amount_minor, currency, merchant_name, category_id
):
    tx_id = str(uuid.uuid4())
    await session.execute(
        sa.text(
            "INSERT INTO transactions "
            "(id, account_id, kind, amount_minor, currency, merchant_name, category_id, status) "
            "VALUES (:id, :account_id, :kind, :amount_minor, :currency, :merchant_name, :category_id, 'pending')"
        ),
        {
            "id": tx_id,
            "account_id": account_id,
            "kind": kind,
            "amount_minor": amount_minor,
            "currency": currency,
            "merchant_name": merchant_name,
            "category_id": category_id,
        },
    )
    await session.commit()
    return tx_id


async def get_transaction(session, tx_id: str):
    result = await session.execute(
        sa.text(
            "SELECT id, account_id, kind, amount_minor, currency, merchant_name, "
            "category_id, status FROM transactions WHERE id = :id"
        ),
        {"id": tx_id},
    )
    return result.first()


async def settle_transaction(session, *, tx, risk_score, use_row_lock: bool):
    """Applies the balance delta and marks the transaction settled. When
    use_row_lock is true (row_lock_contention on a hot account), takes
    SELECT ... FOR UPDATE and holds the transaction open for an extra 150ms
    before committing -- phase-0-base-project-spec.md §7."""
    delta = KIND_SIGN[tx.kind] * tx.amount_minor

    if use_row_lock:
        result = await session.execute(
            sa.text("SELECT balance_minor FROM accounts WHERE id = :id FOR UPDATE"),
            {"id": tx.account_id},
        )
        result.first()
        await session.execute(
            sa.text("UPDATE accounts SET balance_minor = balance_minor + :delta WHERE id = :id"),
            {"delta": delta, "id": tx.account_id},
        )
        await asyncio.sleep(ROW_LOCK_HOLD_SECONDS)
    else:
        await session.execute(
            sa.text("UPDATE accounts SET balance_minor = balance_minor + :delta WHERE id = :id"),
            {"delta": delta, "id": tx.account_id},
        )

    await session.execute(
        sa.text("UPDATE transactions SET status = 'settled', risk_score = :risk_score WHERE id = :id"),
        {"risk_score": risk_score, "id": tx.id},
    )
    await session.commit()

    result = await session.execute(
        sa.text("SELECT balance_minor FROM accounts WHERE id = :id"), {"id": tx.account_id}
    )
    return result.scalar_one()


async def reject_transaction(session, *, tx, risk_score):
    await session.execute(
        sa.text("UPDATE transactions SET status = 'rejected', risk_score = :risk_score WHERE id = :id"),
        {"risk_score": risk_score, "id": tx.id},
    )
    await session.commit()
    result = await session.execute(
        sa.text("SELECT balance_minor FROM accounts WHERE id = :id"), {"id": tx.account_id}
    )
    return result.scalar_one()