From 2d8f528e793996df7ac5143d57baf0394a54b672 Mon Sep 17 00:00:00 2001
From: "pawel.leszczynski" <leszczynski.pawel@gmail.com>
Date: Thu, 14 Mar 2024 13:47:45 +0100
Subject: [PATCH] Spark - prevent NPE if context is null (#2515)

---
 .../io/openlineage/spark/agent/OpenLineageSparkListener.java  | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/integration/spark/app/src/main/java/io/openlineage/spark/agent/OpenLineageSparkListener.java b/integration/spark/app/src/main/java/io/openlineage/spark/agent/OpenLineageSparkListener.java
index d151273b09..b1a3708a36 100644
--- a/integration/spark/app/src/main/java/io/openlineage/spark/agent/OpenLineageSparkListener.java
+++ b/integration/spark/app/src/main/java/io/openlineage/spark/agent/OpenLineageSparkListener.java
@@ -194,7 +194,9 @@ public void onJobEnd(SparkListenerJobEnd jobEnd) {
     ExecutionContext context = rddExecutionRegistry.remove(jobEnd.jobId());
     circuitBreaker.run(
         () -> {
-          context.end(jobEnd);
+          if (context != null) {
+            context.end(jobEnd);
+          }
           return null;
         });
     if (sparkVersion.startsWith("3")) {