"""Readers' chart-data responses carry no SQL.

Superset 6.1.0 drops each query's SQL from chart-data JSON for embedded guest users only
(superset/charts/data/api.py, ChartDataRestApi._send_chart_response). This extends that
condition with the security manager's redacts_query(), when it has one: the Petrolyst managers
return True for readers (petrolyst_reader.ReaderAccessMixin). It runs when the response is
written, so cached results and result_type=query are covered too.

    python3 fix_redact_query_for_readers.py           apply (idempotent)
    python3 fix_redact_query_for_readers.py --verify  exit 1 unless applied
"""

import sys
from pathlib import Path

TARGET = Path("/app/superset/charts/data/api.py")
BEFORE = """\
            if security_manager.is_guest_user():
                for query in queries:
                    query.pop("query", None)
"""
AFTER = """\
            # Petrolyst (patches/fix_redact_query_for_readers.py): readers too.
            if security_manager.is_guest_user() or getattr(
                security_manager, "redacts_query", lambda: False
            )():
                for query in queries:
                    query.pop("query", None)
"""


def main(argv: list[str]) -> int:
    text = TARGET.read_text()
    if "--verify" in argv:
        compile(text, str(TARGET), "exec")
        if text.count(AFTER) == 1 and BEFORE not in text:
            print(f"{TARGET}: readers' SQL redaction in place")
            return 0
        print(f"{TARGET}: readers' SQL redaction missing", file=sys.stderr)
        return 1
    if text.count(AFTER) == 1:
        print(f"{TARGET}: already patched")
        return 0
    if text.count(BEFORE) != 1:
        print(f"{TARGET}: the guest-only redaction was not found once; Superset changed?",
              file=sys.stderr)
        return 1
    TARGET.write_text(text.replace(BEFORE, AFTER))
    print(f"{TARGET}: patched")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))