from ematix_flow import ematix, ManagedTable, Annotated, BigInt, Text, pk

@ematix.connection
class warehouse:
    kind = "postgres"
    url = "${WAREHOUSE_URL}"

class OrdersExtracted(ManagedTable):
    __schema__ = "analytics"; __tablename__ = "orders_extracted"
    order_id: Annotated[BigInt, pk()]
    customer_id: BigInt
    amount_cents: BigInt

class OrdersEnriched(ManagedTable):
    __schema__ = "analytics"; __tablename__ = "orders_enriched"
    order_id: Annotated[BigInt, pk()]
    amount_bucket: Text

class OrdersByCustomer(ManagedTable):
    __schema__ = "analytics"; __tablename__ = "orders_by_customer"
    customer_id: Annotated[BigInt, pk()]
    order_count: BigInt

# Jobs declare their DAG position with depends_on.
@ematix.job(name="extract_orders",
            target=OrdersExtracted, target_connection="warehouse",
            mode="merge", keys=("order_id",))
def extract_orders(conn):
    return "SELECT order_id, customer_id, amount_cents FROM raw.orders"

@ematix.job(name="enrich_orders",
            target=OrdersEnriched, target_connection="warehouse",
            mode="merge", keys=("order_id",),
            depends_on=["extract_orders"])
def enrich_orders(conn):
    return "SELECT order_id, CASE WHEN amount_cents < 10000 THEN 'small' ELSE 'large' END AS amount_bucket FROM analytics.orders_extracted"

@ematix.job(name="aggregate_orders",
            target=OrdersByCustomer, target_connection="warehouse",
            mode="merge", keys=("customer_id",),
            depends_on=["extract_orders"])
def aggregate_orders(conn):
    return "SELECT customer_id, COUNT(*) AS order_count FROM analytics.orders_extracted GROUP BY 1"

# Workflow declares the trigger — fires when upstream_workflow has
# succeeded AND the 21:00 NY tick has reached, since last self-run.
ematix.workflow(
    name="orders_etl",
    triggered_by=["upstream_workflow"],
    schedule="0 21 * * *",
    timezone="America/New_York",
    jobs=["extract_orders", "enrich_orders", "aggregate_orders"],
)