diff --git a/src/datacustomcode/einstein_predictions/spark_default.py b/src/datacustomcode/einstein_predictions/spark_default.py index b6bf7b4..0e43abd 100644 --- a/src/datacustomcode/einstein_predictions/spark_default.py +++ b/src/datacustomcode/einstein_predictions/spark_default.py @@ -39,6 +39,9 @@ _STATUS_SUCCESS = "SUCCESS" _STATUS_ERROR = "ERROR" +# HTTP status considered a successful prediction call. +_HTTP_OK = 200 + class DefaultSparkEinsteinPredictions(SparkEinsteinPredictions): @@ -104,6 +107,8 @@ def einstein_predict_col( def _predict(values_row: Any) -> Dict[str, Optional[str]]: if values_row is None: + # An entirely null features struct is not the normal per-feature null + # case; surface it directly rather than masking it (debuggability). return { "status": _STATUS_ERROR, "response": None, @@ -181,6 +186,21 @@ def _call_predictions( return predictions.predict(request) +def _null_feature_name(features: Dict[str, Any]) -> Optional[str]: + """Return the name of the first null feature value, or ``None``.""" + for name, value in features.items(): + if value is None: + return name + return None + + +def _null_feature_message(name: str) -> str: + return ( + f"Feature '{name}' has null value. Use coalesce() or when() to handle " + f"nulls before calling einstein_predict." + ) + + def _invoke_predictions( predictions: "EinsteinPredictions", model_api_name: str, @@ -190,18 +210,40 @@ def _invoke_predictions( ) -> Dict[str, Any]: from datacustomcode.einstein_predictions.errors import EinsteinPredictionsCallError - response = _call_predictions( - predictions, model_api_name, prediction_type, features, settings - ) - if not response.is_success: - error_code = _extract_error_code(response) + null_feature = _null_feature_name(features) + if null_feature is not None: + message = _null_feature_message(null_feature) + raise EinsteinPredictionsCallError( + f"Einstein Predictions call failed: {message}", + status=None, + error_code=None, + error_message=message, + ) + + try: + response = _call_predictions( + predictions, model_api_name, prediction_type, features, settings + ) + except EinsteinPredictionsCallError: + raise + except Exception as exc: + # Transport/build failures: surface the real error (no masking) so local + # runs stay debuggable. error_code stays None since there is no HTTP status. + raise EinsteinPredictionsCallError( + f"Einstein Predictions call failed: {exc}", + status=None, + error_code=None, + error_message=str(exc), + ) from exc + + if response.status_code != _HTTP_OK: + error_message = json.dumps(response.data) if response.data is not None else None raise EinsteinPredictionsCallError( f"Einstein Predictions call failed: " - f"status_code={response.status_code}, " - f"error_code={error_code!r}, message={response.data!r}", + f"status_code={response.status_code}, message={error_message!r}", status=response.status_code, - error_code=error_code, - error_message=str(response.data) if response.data else None, + error_code=str(response.status_code), + error_message=error_message, ) return response.data or {} @@ -213,27 +255,46 @@ def _invoke_predictions_as_struct( features: Dict[str, Any], settings: Optional[Dict[str, Any]], ) -> Dict[str, Optional[str]]: - response = _call_predictions( - predictions, model_api_name, prediction_type, features, settings - ) - if not response.is_success: + # (a) Customer-actionable data condition — surface the actionable message directly. + null_feature = _null_feature_name(features) + if null_feature is not None: return { "status": _STATUS_ERROR, "response": None, - "error_code": _extract_error_code(response), - "error_message": str(response.data) if response.data else None, + "error_code": None, + "error_message": _null_feature_message(null_feature), } - return { - "status": _STATUS_SUCCESS, - "response": json.dumps(response.data) if response.data is not None else None, - "error_code": None, - "error_message": None, - } + # (b) Transport/build failures — surface the real error (no masking) so local + # runs stay debuggable. error_code stays None since there is no HTTP status. + try: + response = _call_predictions( + predictions, model_api_name, prediction_type, features, settings + ) + except Exception as exc: + return { + "status": _STATUS_ERROR, + "response": None, + "error_code": None, + "error_message": str(exc), + } -def _extract_error_code(response: "PredictionResponse") -> Optional[str]: - if response.data: - error_code = response.data.get("errorCode") - if error_code is not None: - return str(error_code) - return None + if response.status_code == _HTTP_OK: + return { + "status": _STATUS_SUCCESS, + "response": ( + json.dumps(response.data) if response.data is not None else None + ), + "error_code": None, + "error_message": None, + } + + # (c) Non-200 SFAP HTTP error: error_code = status code, error_message = data JSON. + return { + "status": _STATUS_ERROR, + "response": None, + "error_code": str(response.status_code), + "error_message": ( + json.dumps(response.data) if response.data is not None else None + ), + } diff --git a/tests/test_spark_einstein_predictions.py b/tests/test_spark_einstein_predictions.py index b88ef10..d5ec199 100644 --- a/tests/test_spark_einstein_predictions.py +++ b/tests/test_spark_einstein_predictions.py @@ -206,6 +206,7 @@ def test_udf_returns_error_struct_for_null_row(self, mock_struct, mock_udf): out = udf_fn(None) assert out["status"] == _STATUS_ERROR assert out["response"] is None + assert out["error_code"] is None assert "null" in out["error_message"].lower() mock_inner.predict.assert_not_called() @@ -213,7 +214,8 @@ def test_udf_returns_error_struct_for_null_row(self, mock_struct, mock_udf): @patch("pyspark.sql.functions.struct") def test_udf_returns_error_struct_on_http_error(self, mock_struct, mock_udf): """Per-row errors are returned as ``status="ERROR"`` structs so one bad - row does not abort the Spark job.""" + row does not abort the Spark job. ``error_code`` is the HTTP status code + and ``error_message`` is the SFAP body as JSON.""" mock_struct.return_value = MagicMock() mock_udf.return_value = MagicMock() mock_inner = MagicMock() @@ -233,8 +235,99 @@ def test_udf_returns_error_struct_on_http_error(self, mock_struct, mock_udf): assert out["status"] == _STATUS_ERROR assert out["response"] is None - assert out["error_code"] == "UNAVAILABLE" - assert out["error_message"] is not None + assert out["error_code"] == "503" + assert json.loads(out["error_message"]) == {"errorCode": "UNAVAILABLE"} + + @patch("pyspark.sql.functions.udf") + @patch("pyspark.sql.functions.struct") + def test_udf_returns_specific_error_for_null_feature(self, mock_struct, mock_udf): + """A null feature value is a customer-actionable data condition: it is + surfaced with error_code None and an actionable message, never coerced + to the string "None".""" + mock_struct.return_value = MagicMock() + mock_udf.return_value = MagicMock() + mock_inner = MagicMock() + predictions = DefaultSparkEinsteinPredictions(einstein_predictions=mock_inner) + + predictions.einstein_predict_col( + "model1", PredictionType.REGRESSION, {"beds": MagicMock()} + ) + + udf_fn = mock_udf.call_args.args[0] + row = MagicMock() + row.asDict.return_value = {"beds": None} + out = udf_fn(row) + + assert out["status"] == _STATUS_ERROR + assert out["response"] is None + assert out["error_code"] is None + assert "beds" in out["error_message"] + assert "coalesce" in out["error_message"] + mock_inner.predict.assert_not_called() + + @patch("pyspark.sql.functions.udf") + @patch("pyspark.sql.functions.struct") + def test_udf_returns_generic_error_on_transport_failure( + self, mock_struct, mock_udf + ): + """Transport/build exceptions are logged and surfaced with error_code None + and the exception text as error_message so local runs stay debuggable.""" + mock_struct.return_value = MagicMock() + mock_udf.return_value = MagicMock() + mock_inner = MagicMock() + mock_inner.predict.side_effect = RuntimeError("connection refused to 10.0.0.1") + predictions = DefaultSparkEinsteinPredictions(einstein_predictions=mock_inner) + + predictions.einstein_predict_col( + "model1", PredictionType.REGRESSION, {"beds": MagicMock()} + ) + + udf_fn = mock_udf.call_args.args[0] + row = MagicMock() + row.asDict.return_value = {"beds": 3.0} + out = udf_fn(row) + + assert out["status"] == _STATUS_ERROR + assert out["response"] is None + assert out["error_code"] is None + assert out["error_message"] == "connection refused to 10.0.0.1" + + @patch("pyspark.sql.functions.udf") + @patch("pyspark.sql.functions.struct") + def test_udf_passes_through_prediction_failure_as_success( + self, mock_struct, mock_udf + ): + """A 200 response carrying a PredictionFailure stays SUCCESS; the failure + is passed through in ``response`` for the script to handle.""" + mock_struct.return_value = MagicMock() + mock_udf.return_value = MagicMock() + failure_body = { + "results": [ + { + "type": "PredictionFailure", + "error": { + "message": "no match", + "predictionErrorCode": "PREDICTION_ERROR_CODE_NO_MATCH", + }, + } + ] + } + mock_inner = MagicMock() + mock_inner.predict.return_value = _success_response(failure_body) + predictions = DefaultSparkEinsteinPredictions(einstein_predictions=mock_inner) + + predictions.einstein_predict_col( + "model1", PredictionType.BINARY_CLASSIFICATION, {"beds": MagicMock()} + ) + + udf_fn = mock_udf.call_args.args[0] + row = MagicMock() + row.asDict.return_value = {"beds": 3.0} + out = udf_fn(row) + + assert out["status"] == _STATUS_SUCCESS + assert json.loads(out["response"]) == failure_body + assert out["error_code"] is None class TestInvokePredictions: @@ -260,9 +353,37 @@ def test_raises_call_error_on_error_response(self): ) assert excinfo.value.status == 503 - assert excinfo.value.error_code == "UNAVAILABLE" + assert excinfo.value.error_code == "503" + assert excinfo.value.error_message == json.dumps({"errorCode": "UNAVAILABLE"}) assert "503" in str(excinfo.value) - assert "UNAVAILABLE" in str(excinfo.value) + + def test_raises_specific_error_on_null_feature(self): + mock_inner = MagicMock() + + with pytest.raises(EinsteinPredictionsCallError) as excinfo: + _invoke_predictions( + mock_inner, "model", PredictionType.REGRESSION, {"x": None}, None + ) + + assert excinfo.value.status is None + assert excinfo.value.error_code is None + assert "x" in str(excinfo.value.error_message) + assert "coalesce" in str(excinfo.value.error_message) + mock_inner.predict.assert_not_called() + + def test_raises_generic_error_on_transport_failure(self): + mock_inner = MagicMock() + mock_inner.predict.side_effect = RuntimeError("connection refused to 10.0.0.1") + + with pytest.raises(EinsteinPredictionsCallError) as excinfo: + _invoke_predictions( + mock_inner, "model", PredictionType.REGRESSION, {"x": 1.0}, None + ) + + assert excinfo.value.status is None + assert excinfo.value.error_code is None + assert excinfo.value.error_message == "connection refused to 10.0.0.1" + assert "connection refused" in str(excinfo.value) class TestInvokePredictionsAsStruct: @@ -295,8 +416,35 @@ def test_error_returns_error_struct_without_raising(self): assert out["status"] == _STATUS_ERROR assert out["response"] is None - assert out["error_code"] == "UNAVAILABLE" - assert out["error_message"] is not None + assert out["error_code"] == "503" + assert json.loads(out["error_message"]) == {"errorCode": "UNAVAILABLE"} + + def test_null_feature_returns_specific_error_struct(self): + mock_inner = MagicMock() + + out = _invoke_predictions_as_struct( + mock_inner, "model", PredictionType.REGRESSION, {"x": None}, None + ) + + assert out["status"] == _STATUS_ERROR + assert out["response"] is None + assert out["error_code"] is None + assert "x" in out["error_message"] + assert "None" != out["error_message"] + mock_inner.predict.assert_not_called() + + def test_transport_failure_returns_generic_error_struct(self): + mock_inner = MagicMock() + mock_inner.predict.side_effect = RuntimeError("connection refused to 10.0.0.1") + + out = _invoke_predictions_as_struct( + mock_inner, "model", PredictionType.REGRESSION, {"x": 1.0}, None + ) + + assert out["status"] == _STATUS_ERROR + assert out["response"] is None + assert out["error_code"] is None + assert out["error_message"] == "connection refused to 10.0.0.1" class TestDefaultSparkEinsteinPredictionsErrorHandling: @@ -316,4 +464,4 @@ def test_raises_on_error_response(self): ) assert excinfo.value.status == 429 - assert excinfo.value.error_code == "RATE_LIMITED" + assert excinfo.value.error_code == "429"