[Spark] always catch IllegalStateException when accessing SparkSession

CHANGELOG.md

@@ -8,6 +8,8 @@
 ### Fixed
 * **Flink: disable module metadata generation** [`#2507`](https://github.com/OpenLineage/OpenLineage/pull/2531) [@HuangZhenQiu](https://github.com/HuangZhenQiu)
   *Fixes the issue flink module generated contains the internal libs that are not published*
+* **Spark: fix access to active Spark session ** [`#2535`](https://github.com/OpenLineage/OpenLineage/pull/2535) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)  
+  *Always catch `IllegalStateException` when accessing `SparkSession`.*
 
 ## [1.10.2](https://github.com/OpenLineage/OpenLineage/compare/1.9.1...1.10.2) - 2024-03-15
 

integration/spark/shared/src/main/java/io/openlineage/spark/agent/filters/EventFilterUtils.java

@@ -5,6 +5,7 @@
 
 package io.openlineage.spark.agent.filters;
 
+import io.openlineage.spark.agent.util.SparkSessionUtils;
 import io.openlineage.spark.api.OpenLineageContext;
 import java.util.Arrays;
 import java.util.Optional;
@@ -51,7 +52,7 @@ static Optional<LogicalPlan> getLogicalPlan(OpenLineageContext context) {
    * @return
    */
   static boolean isDeltaPlan() {
-    return Optional.of(SparkSession.active())
+    return SparkSessionUtils.activeSession()
         .map(SparkSession::sparkContext)
         .filter(context -> context != null)
         .map(SparkContext::conf)

integration/spark/shared/src/main/java/io/openlineage/spark/agent/lifecycle/plan/BigQueryNodeOutputVisitor.java

@@ -12,6 +12,7 @@
 import com.google.cloud.spark.bigquery.SparkBigQueryConfig;
 import io.openlineage.client.OpenLineage;
 import io.openlineage.spark.agent.util.ReflectionUtils;
+import io.openlineage.spark.agent.util.SparkSessionUtils;
 import io.openlineage.spark.api.DatasetFactory;
 import io.openlineage.spark.api.OpenLineageContext;
 import io.openlineage.spark.api.QueryPlanVisitor;
@@ -54,8 +55,9 @@ public boolean isDefinedAt(LogicalPlan plan) {
         && ((SaveIntoDataSourceCommand) plan).dataSource() instanceof BigQueryRelationProvider;
   }
 
-  private String getFromSaveIntoDataSourceCommand(SaveIntoDataSourceCommand saveCommand) {
-    SQLContext sqlContext = SparkSession.active().sqlContext();
+  private String getFromSaveIntoDataSourceCommand(
+      SaveIntoDataSourceCommand saveCommand, SparkSession session) {
+    SQLContext sqlContext = session.sqlContext();
     BigQueryRelationProvider bqRelationProvider =
         (BigQueryRelationProvider) saveCommand.dataSource();
     SparkBigQueryConfig config =
@@ -88,10 +90,15 @@ private Optional<String> getBigQueryTableName(SparkBigQueryConfig config) {
   @Override
   public List<OpenLineage.OutputDataset> apply(LogicalPlan plan) {
     SaveIntoDataSourceCommand saveCommand = (SaveIntoDataSourceCommand) plan;
-    return Collections.singletonList(
-        factory.getDataset(
-            getFromSaveIntoDataSourceCommand(saveCommand),
-            BIGQUERY_NAMESPACE,
-            saveCommand.schema()));
+    Optional<SparkSession> session = SparkSessionUtils.activeSession();
+    if (session.isPresent()) {
+      return Collections.singletonList(
+          factory.getDataset(
+              getFromSaveIntoDataSourceCommand(saveCommand, session.get()),
+              BIGQUERY_NAMESPACE,
+              saveCommand.schema()));
+    } else {
+      return Collections.emptyList();
+    }
   }
 }

integration/spark/shared/src/main/java/io/openlineage/spark/agent/util/PathUtils.java

@@ -90,7 +90,6 @@ public static DatasetIdentifier fromCatalogTable(
     }
   }
 
-  @SneakyThrows
   private static URI prepareUriFromDefaultTablePath(CatalogTable catalogTable) {
     URI uri =
         SparkSession.active().sessionState().catalog().defaultTablePath(catalogTable.identifier());

integration/spark/shared/src/main/java/io/openlineage/spark/agent/util/SparkSessionUtils.java

@@ -0,0 +1,23 @@
+/*
+/* Copyright 2018-2024 contributors to the OpenLineage project
+/* SPDX-License-Identifier: Apache-2.0
+*/
+
+package io.openlineage.spark.agent.util;
+
+import java.util.Optional;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.spark.sql.SparkSession;
+
+@Slf4j
+public class SparkSessionUtils {
+
+  public static Optional<SparkSession> activeSession() {
+    try {
+      return Optional.of(SparkSession.active());
+    } catch (IllegalStateException e) {
+      log.warn("Cannot obtain active spark session", e);
+      return Optional.empty();
+    }
+  }
+}

integration/spark/shared/src/main/java/io/openlineage/spark/agent/util/SparkVersionUtils.java

@@ -9,6 +9,9 @@
 
 public class SparkVersionUtils {
   public static boolean isSpark3() {
-    return SparkSession.active().version().startsWith("3");
+    return SparkSessionUtils.activeSession()
+        .map(SparkSession::version)
+        .filter(v -> v.startsWith("3"))
+        .isPresent();
   }
 }