SQL & Technical Screen

-- One row per registered account. No prompts, no content, ever.
CREATE TABLE users (
    user_id       BIGINT PRIMARY KEY,
    signed_up_at  TIMESTAMPTZ NOT NULL,
    signup_source TEXT,          -- 'organic','referral','api_docs', ...
    country_bucket TEXT          -- coarse region, never precise geo
);

-- One row per subscription *state change*, Stripe-sourced.
-- A user can appear multiple times (upgrade, downgrade, churn, win-back).
CREATE TABLE subscriptions (
    subscription_id BIGINT PRIMARY KEY,
    user_id         BIGINT REFERENCES users(user_id),
    plan            TEXT NOT NULL,       -- 'free','pro','pro_plus','max'
    mrr             NUMERIC(10,2) NOT NULL,  -- monthly recurring revenue, USD
    started_at      TIMESTAMPTZ NOT NULL,
    canceled_at     TIMESTAMPTZ          -- NULL = still active
);

-- One row per API inference request. METADATA ONLY.
-- No prompt, no response, no token *content* -- just counts and timings.
CREATE TABLE api_requests (
    request_id  BIGINT PRIMARY KEY,
    api_key_id  BIGINT,          -- pseudonymous key, not a user trail
    user_id     BIGINT,          -- nullable: anonymous-tier calls have none
    ts          TIMESTAMPTZ NOT NULL,
    model       TEXT NOT NULL,   -- 'uncensored-model','llama-3.3-70b', ...
    tokens_in   INTEGER NOT NULL,
    tokens_out  INTEGER NOT NULL,
    latency_ms  INTEGER NOT NULL,
    status      TEXT NOT NULL    -- 'ok','rate_limited','error'
);

-- Content-free product events (PostHog-style, person_profiles: identified_only).
-- event names only -- never what the user typed or generated.
CREATE TABLE events (
    event_id   BIGINT PRIMARY KEY,
    user_id    BIGINT,          -- nullable for logged-out / anonymized events
    ts         TIMESTAMPTZ NOT NULL,
    event_name TEXT NOT NULL    -- 'signed_in','chat_created','image_generated',
                                -- 'model_switched','upgrade_clicked'
);

-- On-chain staking/credit activity. Fully PUBLIC data (Ethereum L2), pseudonymous wallets.
-- Join to users ONLY via consented_wallets (a first-party opt-in mapping).
CREATE TABLE onchain_stakes (
    tx_hash     TEXT PRIMARY KEY,
    wallet      TEXT NOT NULL,   -- pseudonymous address, never a user_id
    action      TEXT NOT NULL,   -- 'stake','unstake','mint_credit','burn'
    token_amount  NUMERIC(38,18) NOT NULL,
    block_time  TIMESTAMPTZ NOT NULL
);

-- The ONLY sanctioned wallet<->user bridge: explicit, consented, opt-in.
CREATE TABLE consented_wallets (
    user_id     BIGINT REFERENCES users(user_id),
    wallet      TEXT NOT NULL,
    consented_at TIMESTAMPTZ NOT NULL,
    PRIMARY KEY (user_id, wallet)
);

-- Grain: one row per user per month with their active MRR that month.
-- (Assume a helper that expands subscription state-changes into monthly
--  snapshots; in dbt this is a model, not a subquery you rewrite each time.)
WITH monthly_mrr AS (
    SELECT
        user_id,
        date_trunc('month', month_start) AS month,
        SUM(mrr) AS mrr
    FROM user_month_subscription_snapshot   -- built upstream
    GROUP BY user_id, date_trunc('month', month_start)
),
with_prev AS (
    SELECT
        user_id,
        month,
        mrr,
        LAG(mrr) OVER (
            PARTITION BY user_id ORDER BY month
        ) AS prev_mrr
    FROM monthly_mrr
)
SELECT
    month,
    SUM(CASE WHEN prev_mrr IS NULL AND mrr > 0
             THEN mrr END)                       AS new_mrr,
    SUM(CASE WHEN prev_mrr > 0 AND mrr > prev_mrr
             THEN mrr - prev_mrr END)            AS expansion_mrr,
    SUM(CASE WHEN prev_mrr > 0 AND mrr < prev_mrr AND mrr > 0
             THEN mrr - prev_mrr END)            AS contraction_mrr,  -- negative
    SUM(CASE WHEN prev_mrr > 0 AND mrr = 0
             THEN -prev_mrr END)                 AS churned_mrr       -- negative
FROM with_prev
GROUP BY month
ORDER BY month;

WITH paid_months AS (
    SELECT DISTINCT
        user_id,
        date_trunc('month', month_start) AS month
    FROM user_month_subscription_snapshot
    WHERE mrr > 0
),
cohorts AS (
    SELECT
        user_id,
        month,
        MIN(month) OVER (PARTITION BY user_id) AS cohort_month
    FROM paid_months
)
SELECT
    cohort_month,
    (EXTRACT(YEAR FROM age(month, cohort_month)) * 12
       + EXTRACT(MONTH FROM age(month, cohort_month)))::INT AS months_since_first,
    COUNT(DISTINCT user_id) AS retained_users
FROM cohorts
GROUP BY cohort_month, months_since_first
HAVING COUNT(DISTINCT user_id) >= 20   -- k-anonymity: suppress tiny cohorts
ORDER BY cohort_month, months_since_first;

-- ROW_NUMBER : arbitrary tiebreak, exactly one row per rank -> use for "the Nth".
-- RANK       : ties share a rank, then it SKIPS (1,1,3) -> "2nd" may not exist.
-- DENSE_RANK : ties share a rank, no skip (1,1,2)       -> "2nd distinct value".
WITH spend AS (
    SELECT
        s.plan,
        s.user_id,
        SUM(s.mrr) AS lifetime_spend
    FROM subscriptions s
    GROUP BY s.plan, s.user_id
),
ranked AS (
    SELECT
        plan,
        user_id,
        lifetime_spend,
        ROW_NUMBER() OVER (
            PARTITION BY plan
            ORDER BY lifetime_spend DESC, user_id   -- deterministic tiebreak
        ) AS rn
    FROM spend
)
SELECT plan, user_id, lifetime_spend
FROM ranked
WHERE rn = 2;

WITH signups AS (
    SELECT user_id, signed_up_at
    FROM users
    WHERE signed_up_at >= DATE '2026-07-01'
),
first_chat AS (
    SELECT user_id, MIN(ts) AS first_chat_at
    FROM events
    WHERE event_name = 'chat_created'
    GROUP BY user_id
),
first_paid_inference AS (
    -- first API request made by a user on a paid subscription
    SELECT r.user_id, MIN(r.ts) AS first_paid_at
    FROM api_requests r
    WHERE r.user_id IS NOT NULL
      AND EXISTS (
          SELECT 1 FROM subscriptions s
          WHERE s.user_id = r.user_id
            AND s.mrr > 0
            AND s.started_at <= r.ts
            AND (s.canceled_at IS NULL OR s.canceled_at > r.ts)
      )
    GROUP BY r.user_id
)
SELECT
    COUNT(*)                                              AS signed_up,
    COUNT(fc.first_chat_at)                               AS reached_first_chat,
    COUNT(fp.first_paid_at)                               AS reached_first_paid,
    ROUND(100.0 * COUNT(fc.first_chat_at) / COUNT(*), 1)  AS pct_to_chat,
    ROUND(100.0 * COUNT(fp.first_paid_at)
          / NULLIF(COUNT(fc.first_chat_at), 0), 1)        AS pct_chat_to_paid
FROM signups su
LEFT JOIN first_chat fc            USING (user_id)
LEFT JOIN first_paid_inference fp  USING (user_id)
WHERE su.user_id IN (SELECT user_id FROM signups);

WITH activity AS (
    SELECT DISTINCT user_id, date_trunc('day', ts)::date AS active_day
    FROM events
    WHERE user_id IS NOT NULL
),
first_seen AS (
    SELECT user_id, MIN(active_day) AS day0
    FROM activity
    GROUP BY user_id
)
SELECT
    a.active_day - f.day0                       AS day_n,
    COUNT(DISTINCT a.user_id)                   AS retained
FROM activity a
JOIN first_seen f USING (user_id)
WHERE a.active_day - f.day0 BETWEEN 0 AND 30
GROUP BY a.active_day - f.day0
HAVING COUNT(DISTINCT a.user_id) >= 20          -- suppress small cells
ORDER BY day_n;

-- Per-model usage, but only report a model if >= k distinct users touched it.
SELECT
    model,
    COUNT(*)                    AS requests,
    SUM(tokens_in + tokens_out) AS total_tokens,
    ROUND(AVG(latency_ms))      AS avg_latency_ms
FROM api_requests
WHERE ts >= DATE '2026-08-01'
GROUP BY model
HAVING COUNT(DISTINCT user_id) >= 20   -- k = 20; drop long-tail single-user models
ORDER BY requests DESC;

-- WRONG: window functions are not allowed in WHERE.
-- SELECT * FROM api_requests
-- WHERE ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY ts DESC) = 1;  -- error

-- RIGHT: compute in a CTE, filter outside.
WITH ranked AS (
    SELECT
        r.*,
        ROW_NUMBER() OVER (PARTITION BY user_id ORDER BY ts DESC) AS rn
    FROM api_requests r
    WHERE user_id IS NOT NULL
)
SELECT * FROM ranked WHERE rn = 1;   -- (QUALIFY rn = 1 in BigQuery/Snowflake)

-- DANGER: if ANY api_requests.user_id is NULL, this returns NOTHING.
-- SELECT * FROM users
-- WHERE user_id NOT IN (SELECT user_id FROM api_requests);

-- SAFE: NOT EXISTS is null-proof and usually faster.
SELECT u.*
FROM users u
WHERE NOT EXISTS (
    SELECT 1 FROM api_requests r WHERE r.user_id = u.user_id
);

-- Keep one row per logical event, preferring the earliest ingest.
WITH deduped AS (
    SELECT
        e.*,
        ROW_NUMBER() OVER (
            PARTITION BY event_id          -- the true natural key
            ORDER BY ts                    -- keep the first-seen copy
        ) AS rn
    FROM events e
)
SELECT * FROM deduped WHERE rn = 1;

-- Align usage to the SAME 00:00 UTC boundary the staking allocation resets on.
SELECT
    date_trunc('day', ts AT TIME ZONE 'UTC') AS utc_day,
    COUNT(*)                                 AS requests,
    SUM(tokens_out)                          AS tokens_served
FROM api_requests
GROUP BY date_trunc('day', ts AT TIME ZONE 'UTC')
ORDER BY utc_day;
-- Half-open intervals [start, end) beat BETWEEN for timestamps:
--   ts >= '2026-08-01' AND ts < '2026-09-01'
-- avoids the "misses/double-counts the last millisecond" BETWEEN bug.

WITH flagged AS (
    SELECT
        user_id,
        ts,
        CASE WHEN ts - LAG(ts) OVER (PARTITION BY user_id ORDER BY ts)
                  > INTERVAL '30 minutes'
             OR LAG(ts) OVER (PARTITION BY user_id ORDER BY ts) IS NULL
             THEN 1 ELSE 0 END AS is_new_session
    FROM events
    WHERE user_id IS NOT NULL
),
sessioned AS (
    SELECT
        user_id,
        ts,
        SUM(is_new_session) OVER (
            PARTITION BY user_id ORDER BY ts
            ROWS UNBOUNDED PRECEDING
        ) AS session_seq
    FROM flagged
)
SELECT
    user_id, session_seq,
    MIN(ts) AS session_start,
    MAX(ts) AS session_end,
    COUNT(*) AS events_in_session
FROM sessioned
GROUP BY user_id, session_seq;

"""Aggregate content-free API-usage rows into k-anonymous daily rollups.

Pure function, no I/O -- easy to unit-test and safe to reuse in a dbt
Python model, an Airflow task, or an ad-hoc notebook.
"""
from __future__ import annotations

import pandas as pd


def daily_model_rollup(
    requests: pd.DataFrame,
    *,
    k: int = 20,
    tz: str = "UTC",  # align to the 00:00 UTC staking-reset boundary
) -> pd.DataFrame:
    """Roll usage metadata up to (utc_day, model), suppressing small cells.

    Args:
        requests: rows with columns
            [ts, user_id, model, tokens_in, tokens_out, latency_ms].
            Carries METADATA ONLY -- never prompt/response content.
        k: minimum distinct users required to emit a (day, model) cell.
        tz: timezone to bucket ``ts`` into before truncating to the day.

    Returns:
        One row per (utc_day, model) with >= k distinct users.
    """
    required = {"ts", "user_id", "model", "tokens_in", "tokens_out", "latency_ms"}
    missing = required - set(requests.columns)
    if missing:
        raise ValueError(f"missing required columns: {sorted(missing)}")

    df = requests.copy()
    df["utc_day"] = (
        pd.to_datetime(df["ts"], utc=True).dt.tz_convert(tz).dt.floor("D")
    )
    df["tokens_total"] = df["tokens_in"] + df["tokens_out"]

    grouped = df.groupby(["utc_day", "model"], as_index=False).agg(
        requests=("ts", "size"),
        distinct_users=("user_id", "nunique"),
        tokens_total=("tokens_total", "sum"),
        avg_latency_ms=("latency_ms", "mean"),
    )
    # k-anonymity: never emit a cell backed by fewer than k users.
    return grouped[grouped["distinct_users"] >= k].reset_index(drop=True)

import pandas as pd
import pytest

from rollup import daily_model_rollup


def _rows(n_users: int, model: str = "uncensored-model") -> pd.DataFrame:
    return pd.DataFrame({
        "ts": ["2026-08-01T12:00:00Z"] * n_users,
        "user_id": range(n_users),
        "model": [model] * n_users,
        "tokens_in": [10] * n_users,
        "tokens_out": [20] * n_users,
        "latency_ms": [100] * n_users,
    })


def test_suppresses_cells_below_k():
    out = daily_model_rollup(_rows(5), k=20)
    assert out.empty  # 5 users < k -> nothing emitted


def test_emits_cell_at_or_above_k():
    out = daily_model_rollup(_rows(25), k=20)
    assert len(out) == 1
    assert out.loc[0, "distinct_users"] == 25
    assert out.loc[0, "tokens_total"] == 25 * 30


def test_rejects_missing_columns():
    with pytest.raises(ValueError, match="missing required columns"):
        daily_model_rollup(pd.DataFrame({"ts": []}), k=1)