consumer.py

"""
consumer.py  (v1, with v2 enrichment hook)

Reads records from the streaming bus, scores sentiment with VADER, optionally
runs the v2 enrichment stage, and writes enriched records to S3 (or the local
sink) as micro-batched newline-delimited JSON for Snowpipe to auto-ingest.

Usage:
    python consumer.py                     # consume Kinesis -> S3
    LOCAL_MODE=true python consumer.py     # consume local JSONL -> local sink
    ENRICHMENT_ENABLED=true python consumer.py   # turn on v2 enrichment

A micro-batch is flushed whenever MICRO_BATCH_SECONDS elapses OR
MICRO_BATCH_MAX_RECORDS records have accumulated, whichever comes first. A
background thread checks the timer once a second so a quiet stream still
flushes on schedule (GAPS.md A2), and the buffer is also flushed via atexit.

Malformed records (bad JSON, or anything that blows up enrichment) are logged,
counted, and appended to ./rejects.jsonl as a dead-letter -- they no longer
kill the process (GAPS.md B1).

NOTE: This is a single-process shard-iterator consumer suitable for the MVP and
moderate volumes. For production / high throughput use the Kinesis Client
Library (KCL) or a Lambda event source mapping instead. It checkpoints
per-shard sequence numbers to disk so a restart resumes with
AFTER_SEQUENCE_NUMBER instead of losing everything published while it was down
(GAPS.md A1), but this is still best-effort, not exactly-once.
"""
from __future__ import annotations

import argparse
import atexit
import json
import logging
import os
import sys
import threading
import time
import uuid
from datetime import datetime, timezone
from io import StringIO

from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer

import config

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s %(levelname)s [consumer] %(message)s"
)
log = logging.getLogger("consumer")

_analyzer = SentimentIntensityAnalyzer()

# Optional v2 enrichment. Imported lazily and only used when enabled, so v1 runs
# without langdetect / KeyBERT installed.
_enricher = None
if config.ENRICHMENT_ENABLED:
    try:
        from enrichment import Enricher

        _enricher = Enricher()
        log.info("enrichment stage ENABLED (v2)")
    except Exception as exc:  # pragma: no cover
        log.warning("enrichment requested but could not load: %s", exc)
        _enricher = None

# Reject counter + dead-letter file (GAPS.md B1). A malformed record is logged,
# counted, and appended here instead of crashing the consumer.
_reject_count = 0
REJECT_LOG_PATH = "./rejects.jsonl"


def _log_reject(raw: str, reason: str) -> None:
    """Count a rejected record, log it, and best-effort dead-letter it to disk."""
    global _reject_count
    _reject_count += 1
    log.warning("rejecting malformed record (#%d): %s", _reject_count, reason)
    try:
        with open(REJECT_LOG_PATH, "a", encoding="utf-8") as fh:
            fh.write(
                json.dumps(
                    {
                        "raw": raw,
                        "reason": reason,
                        "rejected_at": datetime.now(timezone.utc).isoformat(),
                    }
                )
                + "\n"
            )
    except Exception as exc:  # dead-letter write is best-effort
        log.debug("could not write to reject log: %s", exc)


def score_sentiment(text: str) -> tuple[str, float]:
    """Return (label, compound_score) using VADER's standard thresholds."""
    compound = _analyzer.polarity_scores(text or "")["compound"]
    if compound >= 0.05:
        label = "positive"
    elif compound <= -0.05:
        label = "negative"
    else:
        label = "neutral"
    return label, round(compound, 4)


# Required keys and the schema contract enforced before a record is admitted
# (GAPS.md S7). With the v3 connectors the bus carries untrusted public data, so
# oversized or wrong-typed fields would otherwise flow into the warehouse and
# hit the pipe's silent ON_ERROR = CONTINUE drop.
REQUIRED_KEYS = ("review_id", "text")
MAX_TEXT_CHARS = 10_000


class RecordValidationError(ValueError):
    """Raised when an inbound record doesn't meet the schema contract."""


def validate_record(record: dict) -> dict:
    """Check required keys, coerce types, and cap text length (GAPS.md S7).

    Returns the (possibly coerced) record; raises RecordValidationError so the
    caller can dead-letter it.
    """
    if not isinstance(record, dict):
        raise RecordValidationError(f"expected a JSON object, got {type(record).__name__}")
    for key in REQUIRED_KEYS:
        if key not in record or record[key] is None:
            raise RecordValidationError(f"missing required field '{key}'")
    if not isinstance(record["text"], str):
        raise RecordValidationError(
            f"'text' must be a string, got {type(record['text']).__name__}"
        )
    if len(record["text"]) > MAX_TEXT_CHARS:
        log.debug("truncating oversized text on %s", record.get("review_id"))
        record["text"] = record["text"][:MAX_TEXT_CHARS]
    star = record.get("star_rating")
    if star is not None and not isinstance(star, (int, float)):
        try:
            record["star_rating"] = int(star)
        except (TypeError, ValueError):
            record["star_rating"] = None
    return record


def enrich_record(record: dict) -> dict | None:
    """Validate, score sentiment, and (optionally) run v2 enrichment.

    Returns the enriched record, or None if it was filtered (e.g. a duplicate).
    Raises RecordValidationError for records that fail the schema contract.
    """
    record = validate_record(record)
    label, score = score_sentiment(record.get("text", ""))
    record["sentiment_label"] = label
    record["sentiment_score"] = score
    record["processed_at"] = datetime.now(timezone.utc).isoformat()

    if _enricher is not None:
        record = _enricher.enrich(record)
        if record is None:  # dropped as duplicate
            return None
    return record


class MicroBatchSink:
    """Buffers enriched records and flushes them to S3 (or the local sink).

    _buf is touched from both the main consumer loop and the idle-flush
    background thread, so all access goes through _lock.
    """

    def __init__(self) -> None:
        self.local = config.LOCAL_MODE
        self._buf: list[dict] = []
        self._last_flush = time.monotonic()
        self._lock = threading.Lock()
        # Distinguishes batches flushed within the same second; without it the
        # timestamp-only key collides and one batch silently overwrites the
        # other (both locally and in S3).
        self._batch_seq = 0
        self._run_id = uuid.uuid4().hex[:8]
        if self.local:
            os.makedirs(config.LOCAL_SINK_DIR, exist_ok=True)
            self._s3 = None
            log.info("LOCAL_MODE on -> micro-batches to %s", config.LOCAL_SINK_DIR)
        else:
            import boto3

            self._s3 = boto3.client("s3", region_name=config.AWS_REGION)
            log.info("micro-batches -> s3://%s/%s", config.S3_BUCKET, config.S3_PREFIX)

    def add(self, record: dict) -> None:
        with self._lock:
            self._buf.append(record)
        self._maybe_flush()

    def _maybe_flush(self, force: bool = False) -> None:
        with self._lock:
            if not self._buf:
                self._last_flush = time.monotonic()
                return
            elapsed = time.monotonic() - self._last_flush
            if (
                force
                or len(self._buf) >= config.MICRO_BATCH_MAX_RECORDS
                or elapsed >= config.MICRO_BATCH_SECONDS
            ):
                self._flush_locked()

    def _flush_locked(self) -> None:
        """Flush self._buf. Caller must hold self._lock."""
        if not self._buf:
            return
        now = datetime.now(timezone.utc)
        # Newline-delimited JSON -- the format Snowpipe's COPY expects.
        body = StringIO()
        for rec in self._buf:
            body.write(json.dumps(rec))
            body.write("\n")
        payload = body.getvalue().encode("utf-8")

        # Hive-style partition path keeps S3 browsable and Snowpipe-friendly.
        # run_id + sequence make the key unique even when several batches flush
        # inside the same second (%H%M%S alone is not enough at high volume).
        self._batch_seq += 1
        key = (
            f"{config.S3_PREFIX}/{now:%Y/%m/%d/%H}/"
            f"batch-{now:%Y%m%dT%H%M%S}-{self._run_id}-"
            f"{self._batch_seq:06d}-{len(self._buf)}.json"
        )
        if self.local:
            path = os.path.join(config.LOCAL_SINK_DIR, key.replace("/", "_"))
            with open(path, "wb") as fh:
                fh.write(payload)
            log.info("flushed %d records -> %s", len(self._buf), path)
        else:
            self._s3.put_object(Bucket=config.S3_BUCKET, Key=key, Body=payload)
            log.info("flushed %d records -> s3://%s/%s",
                     len(self._buf), config.S3_BUCKET, key)

        self._buf.clear()
        self._last_flush = time.monotonic()

    def flush(self) -> None:
        self._maybe_flush(force=True)


# --------------------------------------------------------------------------- #
# Record sources
# --------------------------------------------------------------------------- #
def _iter_local_stream():
    """Tail the local JSONL stream file (LOCAL_MODE)."""
    path = config.LOCAL_STREAM_PATH
    while not os.path.exists(path):
        log.info("waiting for local stream file %s ...", path)
        time.sleep(1)
    with open(path, encoding="utf-8") as fh:
        # start at end so we behave like a 'LATEST' iterator
        fh.seek(0, os.SEEK_END)
        while True:
            line = fh.readline()
            if not line:
                time.sleep(0.25)
                continue
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError as exc:
                _log_reject(line, f"json decode error: {exc}")
                continue
            yield record


# --------------------------------------------------------------------------- #
# Kinesis checkpointing (GAPS.md A1) -- per-shard sequence numbers persisted to
# disk so a restart resumes with AFTER_SEQUENCE_NUMBER instead of LATEST.
# --------------------------------------------------------------------------- #
def _load_checkpoints(path: str) -> dict[str, str]:
    if not os.path.exists(path):
        return {}
    try:
        with open(path, encoding="utf-8") as fh:
            return json.load(fh)
    except Exception as exc:
        log.warning("could not read checkpoint file %s (%s); starting fresh", path, exc)
        return {}


def _save_checkpoints(path: str, checkpoints: dict[str, str]) -> None:
    try:
        tmp = f"{path}.tmp"
        with open(tmp, "w", encoding="utf-8") as fh:
            json.dump(checkpoints, fh)
        os.replace(tmp, path)
    except Exception as exc:  # checkpointing is best-effort; never crash on it
        log.warning("could not write checkpoint file %s: %s", path, exc)


def _iter_kinesis():
    """Poll all shards of the Kinesis stream and yield decoded records.

    Resumes each shard from its last checkpointed sequence number when one
    exists (A1); a shard with no prior checkpoint still starts at LATEST, so
    the "start consumer before producer" rule holds for a first run. Backs off
    on ProvisionedThroughputExceededException, refreshes expired iterators,
    and periodically re-describes the stream to pick up new shards from
    resharding (B2).
    """
    import boto3
    from botocore.exceptions import ClientError

    client = boto3.client("kinesis", region_name=config.AWS_REGION)
    checkpoint_path = config.KINESIS_CHECKPOINT_PATH
    checkpoints = _load_checkpoints(checkpoint_path)

    def _describe_shards():
        return client.describe_stream(StreamName=config.KINESIS_STREAM_NAME)[
            "StreamDescription"
        ]["Shards"]

    def _get_iterator(shard_id: str) -> str:
        checkpoint = checkpoints.get(shard_id)
        if checkpoint:
            return client.get_shard_iterator(
                StreamName=config.KINESIS_STREAM_NAME,
                ShardId=shard_id,
                ShardIteratorType="AFTER_SEQUENCE_NUMBER",
                StartingSequenceNumber=checkpoint,
            )["ShardIterator"]
        return client.get_shard_iterator(
            StreamName=config.KINESIS_STREAM_NAME,
            ShardId=shard_id,
            ShardIteratorType="LATEST",
        )["ShardIterator"]

    iterators: dict[str, str | None] = {}
    for shard in _describe_shards():
        iterators[shard["ShardId"]] = _get_iterator(shard["ShardId"])
    resumed = sum(1 for sid in iterators if sid in checkpoints)
    log.info("polling %d shard(s) (%d resumed from checkpoint)", len(iterators), resumed)

    last_reshard_check = time.monotonic()
    backoff = 1.0

    while True:
        for shard_id, shard_it in list(iterators.items()):
            if shard_it is None:
                continue
            try:
                resp = client.get_records(ShardIterator=shard_it, Limit=500)
            except ClientError as exc:
                code = exc.response.get("Error", {}).get("Code", "")
                if code == "ProvisionedThroughputExceededException":
                    log.warning(
                        "throughput exceeded on shard %s; backing off %.1fs",
                        shard_id, backoff,
                    )
                    time.sleep(backoff)
                    backoff = min(backoff * 2, 30.0)
                    continue
                if code == "ExpiredIteratorException":
                    log.warning("iterator expired for shard %s; refreshing", shard_id)
                    iterators[shard_id] = _get_iterator(shard_id)
                    continue
                log.warning("get_records failed on shard %s (%s); will retry", shard_id, code)
                time.sleep(1.0)
                continue
            backoff = 1.0

            records = resp.get("Records", [])
            for rec in records:
                for line in rec["Data"].decode("utf-8").splitlines():
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        yield json.loads(line)
                    except json.JSONDecodeError as exc:
                        _log_reject(line, f"json decode error: {exc}")
            if records:
                checkpoints[shard_id] = records[-1]["SequenceNumber"]
                _save_checkpoints(checkpoint_path, checkpoints)
            iterators[shard_id] = resp.get("NextShardIterator")

        if time.monotonic() - last_reshard_check > 60:
            last_reshard_check = time.monotonic()
            try:
                current = {s["ShardId"] for s in _describe_shards()}
                for shard_id in current - iterators.keys():
                    log.info("detected new shard %s (resharding); adding", shard_id)
                    iterators[shard_id] = _get_iterator(shard_id)
            except Exception as exc:
                log.warning("resharding check failed: %s", exc)

        time.sleep(1)  # respect Kinesis read throughput limits


def run() -> None:
    if not config.LOCAL_MODE:
        for p in config.validate_for("kinesis") + config.validate_for("s3"):
            log.warning(p)

    sink = MicroBatchSink()
    atexit.register(sink.flush)

    # Idle-flush thread: _maybe_flush() is otherwise only called from add(), so
    # a quiet stream would let records sit in the buffer indefinitely (A2).
    stop_event = threading.Event()

    def _idle_flush_loop() -> None:
        while not stop_event.wait(1.0):
            sink._maybe_flush()

    threading.Thread(target=_idle_flush_loop, daemon=True).start()

    source = _iter_local_stream() if config.LOCAL_MODE else _iter_kinesis()
    processed = 0
    try:
        for record in source:
            try:
                enriched = enrich_record(record)
            except RecordValidationError as exc:
                _log_reject(json.dumps(record, default=str), f"schema violation: {exc}")
                continue
            except Exception as exc:
                _log_reject(json.dumps(record, default=str), f"enrich_record failed: {exc}")
                continue
            if enriched is None:
                continue
            sink.add(enriched)
            processed += 1
            if processed % 100 == 0:
                log.info("processed %d records", processed)
    except KeyboardInterrupt:
        log.info("interrupted by user")
    finally:
        stop_event.set()
        sink.flush()
        log.info("done. total processed: %d, rejected: %d", processed, _reject_count)


def main(argv: list[str] | None = None) -> int:
    argparse.ArgumentParser(description="Kinesis consumer + VADER (v1)").parse_args(argv)
    run()
    return 0


if __name__ == "__main__":
    sys.exit(main())