diff --git a/backend/data_loader.py b/backend/data_loader.py index be1303d..75c4528 100644 --- a/backend/data_loader.py +++ b/backend/data_loader.py @@ -4,6 +4,7 @@ Provides efficient queries with caching. """ import logging +import re import pandas as pd import numpy as np @@ -291,13 +292,28 @@ _GIAS_CODE_COLUMN_NAMES = ( "admissions_policy_code", ) +_MISSING_COLUMN_RE = re.compile(r'column "?(?:s\.)?(\w+)"? does not exist') + + +def _missing_column_name(exc: Exception) -> Optional[str]: + """Name of the missing column from a psycopg2 UndefinedColumn error. + + Inspects exc.orig (the DBAPI error), whose message names only the + offending column — str(exc) also embeds the full SQL statement, which + contains every column name and therefore must not be matched against. + """ + orig = getattr(exc, "orig", None) + match = _MISSING_COLUMN_RE.search(str(orig) if orig is not None else str(exc)) + return match.group(1) if match else None + def load_school_data_as_dataframe() -> pd.DataFrame: """Load all school + KS2 data as a pandas DataFrame.""" try: df = pd.read_sql(_MAIN_QUERY, engine) except sqlalchemy.exc.ProgrammingError as exc: - if any(col in str(exc) for col in _GIAS_CODE_COLUMN_NAMES): + missing = _missing_column_name(exc) + if missing in _GIAS_CODE_COLUMN_NAMES: logging.getLogger(__name__).warning( "marts predate the GIAS code migration — falling back to " "legacy name-column query: %s", @@ -308,7 +324,7 @@ def load_school_data_as_dataframe() -> pd.DataFrame: except Exception as exc2: print(f"Warning: Could not load school data from marts: {exc2}") return pd.DataFrame() - elif "has_sixth_form" in str(exc): + elif missing == "has_sixth_form": logging.getLogger(__name__).warning( "marts.dim_school is missing has_sixth_form (pipeline hasn't " "rebuilt the mart yet on this DB) — retrying without it: %s", diff --git a/backend/tests/test_gias_translation.py b/backend/tests/test_gias_translation.py index 682cff4..414b8e4 100644 --- a/backend/tests/test_gias_translation.py +++ b/backend/tests/test_gias_translation.py @@ -4,7 +4,7 @@ rest of the backend sees must carry today's name strings.""" import numpy as np import pandas as pd -from backend.data_loader import translate_gias_code_columns +from backend.data_loader import _missing_column_name, translate_gias_code_columns from backend.gias_codes import ESTABLISHMENT_STATUS, PHASE_OF_EDUCATION @@ -44,6 +44,39 @@ def test_missing_code_columns_are_a_noop(): assert out.iloc[0]["status"] == "Open" +def _fake_exc(orig_message): + """A stand-in for sqlalchemy.exc.ProgrammingError: str(exc) embeds the + full SQL statement (deliberately containing every column name below, to + prove the matcher doesn't fall back to it), while .orig carries the real + DBAPI error message naming only the offending column.""" + exc = Exception( + "SELECT s.phase_code, s.school_type_code, s.religious_character_code, " + "s.status_code, s.admissions_policy_code, s.has_sixth_form FROM ... " + f"[SQL: ...] (Background on this error at: https://...)" + ) + exc.orig = Exception(orig_message) if orig_message is not None else None + return exc + + +def test_missing_column_name_quoted(): + assert _missing_column_name(_fake_exc('column "phase_code" does not exist')) == "phase_code" + + +def test_missing_column_name_unquoted(): + assert _missing_column_name(_fake_exc("column phase_code does not exist")) == "phase_code" + + +def test_missing_column_name_table_prefixed(): + assert ( + _missing_column_name(_fake_exc("column s.has_sixth_form does not exist")) + == "has_sixth_form" + ) + + +def test_missing_column_name_no_match_returns_none(): + assert _missing_column_name(_fake_exc("relation \"marts.dim_school\" does not exist")) is None + + def test_load_school_data_survives_premigration_marts(monkeypatch): """Real prod state until the nightly pipeline first rebuilds the mart with the GIAS code columns: marts.dim_school still has the old name columns @@ -75,9 +108,12 @@ def test_load_school_data_survives_premigration_marts(monkeypatch): calls.append(query) if len(calls) == 1: raise sqlalchemy.exc.ProgrammingError( - "(psycopg2.errors.UndefinedColumn) column s.phase_code does not exist", - None, - None, + statement=str(data_loader._MAIN_QUERY), + params=None, + orig=Exception( + "(psycopg2.errors.UndefinedColumn) column s.phase_code " + "does not exist\nLINE 5: s.phase_code," + ), ) return good_df.copy() diff --git a/backend/tests/test_sixth_form_flag.py b/backend/tests/test_sixth_form_flag.py index af32799..a5a87f2 100644 --- a/backend/tests/test_sixth_form_flag.py +++ b/backend/tests/test_sixth_form_flag.py @@ -148,10 +148,14 @@ def test_load_school_data_survives_missing_has_sixth_form_column(monkeypatch): def fake_read_sql(query, con): calls.append(query) if len(calls) == 1: + # The statement text still contains phase_code, school_type_code, + # etc. (it's the full _MAIN_QUERY SELECT list) — that's exactly + # the collision this test guards against: matching must be done + # against exc.orig (the DBAPI error), not str(exc)/the statement. raise sqlalchemy.exc.ProgrammingError( - "SELECT ...", - None, - Exception( + statement=str(data_loader._MAIN_QUERY), + params=None, + orig=Exception( "(psycopg2.errors.UndefinedColumn) column s.has_sixth_form " "does not exist" ),