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

@ematix.connection
class app_db:
    kind = "postgres"
    url = "${APP_DB_URL}"

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

# Target tables — typed Python; framework migrates on first run.
class Events(ManagedTable):
    __schema__ = "analytics"; __tablename__ = "events"
    event_id: Annotated[BigInt, pk()]
    name: Text | None
    received_at: TimestampTZ

class DailyEvents(ManagedTable):
    __schema__ = "analytics"; __tablename__ = "events_daily"
    day: Annotated[Text, pk()]
    count: BigInt

# Two jobs — same decorator surface as v0.5.
@ematix.job(
    name="ingest_events",
    source_connection="app_db",
    target=Events,
    target_connection="warehouse",
    schedule="*/5 * * * *",
    mode="append",
)
def ingest_events(conn):
    return "SELECT event_id, name, received_at FROM public.events"

@ematix.job(
    name="rollup_daily",
    source_connection="warehouse",
    target=DailyEvents,
    target_connection="warehouse",
    schedule="0 1 * * *",
    mode="merge", keys=("day",),
)
def rollup_daily(conn):
    return ("SELECT date_trunc('day', received_at)::text AS day, "
            "COUNT(*) AS count FROM analytics.events GROUP BY 1")

# Workflow — names the group + the DAG between its jobs.
ematix.workflow(
    name="events_etl",
    jobs=["ingest_events", "rollup_daily"],
    depends_on={"rollup_daily": ["ingest_events"]},
)