Flink: Added error handling and default logic for Flink version detec…

…tion (apache#9452)
cdouglas · Apr 22, 2024 · 17d3cdf · 17d3cdf
1 parent 98b0ade
commit 17d3cdf
Show file tree

Hide file tree

Showing 6 changed files with 162 additions and 9 deletions.
diff --git a/flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/util/FlinkPackage.java b/flink/v1.16/flink/src/main/java/org/apache/iceberg/flink/util/FlinkPackage.java
@@ -18,16 +18,44 @@
  */
 package org.apache.iceberg.flink.util;
 
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
 
 public class FlinkPackage {
-  /** Choose {@link DataStream} class because it is one of the core Flink API. */
-  private static final String VERSION = DataStream.class.getPackage().getImplementationVersion();
+
+  private static final AtomicReference<String> VERSION = new AtomicReference<>();
+  public static final String FLINK_UNKNOWN_VERSION = "FLINK-UNKNOWN-VERSION";
 
   private FlinkPackage() {}
 
   /** Returns Flink version string like x.y.z */
   public static String version() {
-    return VERSION;
+    if (null == VERSION.get()) {
+      String detectedVersion = null;
+      try {
+        detectedVersion = versionFromJar();
+        // use unknown version in case exact implementation version can't be found from the jar
+        // (this can happen if the DataStream class appears multiple times in the same classpath
+        // such as with shading)
+        detectedVersion = detectedVersion != null ? detectedVersion : FLINK_UNKNOWN_VERSION;
+      } catch (Exception e) {
+        detectedVersion = FLINK_UNKNOWN_VERSION;
+      }
+      VERSION.set(detectedVersion);
+    }
+
+    return VERSION.get();
+  }
+
+  @VisibleForTesting
+  static String versionFromJar() {
+    // Choose {@link DataStream} class because it is one of the core Flink API
+    return DataStream.class.getPackage().getImplementationVersion();
+  }
+
+  @VisibleForTesting
+  static void setVersion(String version) {
+    VERSION.set(version);
   }
 }
diff --git a/flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/util/TestFlinkPackage.java b/flink/v1.16/flink/src/test/java/org/apache/iceberg/flink/util/TestFlinkPackage.java
@@ -20,6 +20,8 @@
 
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
 public class TestFlinkPackage {
 
@@ -28,4 +30,25 @@ public class TestFlinkPackage {
   public void testVersion() {
     Assert.assertEquals("1.16.2", FlinkPackage.version());
   }
+
+  @Test
+  public void testDefaultVersion() {
+    // It's difficult to reproduce a reflection error in a unit test, so we just inject a mocked
+    // fault to test the default logic
+
+    // First make sure we're not caching a version result from a previous test
+    FlinkPackage.setVersion(null);
+    try (MockedStatic<FlinkPackage> mockedStatic = Mockito.mockStatic(FlinkPackage.class)) {
+      mockedStatic.when(FlinkPackage::versionFromJar).thenThrow(RuntimeException.class);
+      mockedStatic.when(FlinkPackage::version).thenCallRealMethod();
+      Assert.assertEquals(FlinkPackage.FLINK_UNKNOWN_VERSION, FlinkPackage.version());
+    }
+    FlinkPackage.setVersion(null);
+    try (MockedStatic<FlinkPackage> mockedStatic = Mockito.mockStatic(FlinkPackage.class)) {
+      mockedStatic.when(FlinkPackage::versionFromJar).thenReturn(null);
+      mockedStatic.when(FlinkPackage::version).thenCallRealMethod();
+      FlinkPackage.setVersion(null);
+      Assert.assertEquals(FlinkPackage.FLINK_UNKNOWN_VERSION, FlinkPackage.version());
+    }
+  }
 }
diff --git a/flink/v1.17/flink/src/main/java/org/apache/iceberg/flink/util/FlinkPackage.java b/flink/v1.17/flink/src/main/java/org/apache/iceberg/flink/util/FlinkPackage.java
@@ -18,16 +18,44 @@
  */
 package org.apache.iceberg.flink.util;
 
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
 
 public class FlinkPackage {
-  /** Choose {@link DataStream} class because it is one of the core Flink API. */
-  private static final String VERSION = DataStream.class.getPackage().getImplementationVersion();
+
+  private static final AtomicReference<String> VERSION = new AtomicReference<>();
+  public static final String FLINK_UNKNOWN_VERSION = "FLINK-UNKNOWN-VERSION";
 
   private FlinkPackage() {}
 
   /** Returns Flink version string like x.y.z */
   public static String version() {
-    return VERSION;
+    if (null == VERSION.get()) {
+      String detectedVersion = null;
+      try {
+        detectedVersion = versionFromJar();
+        // use unknown version in case exact implementation version can't be found from the jar
+        // (this can happen if the DataStream class appears multiple times in the same classpath
+        // such as with shading)
+        detectedVersion = detectedVersion != null ? detectedVersion : FLINK_UNKNOWN_VERSION;
+      } catch (Exception e) {
+        detectedVersion = FLINK_UNKNOWN_VERSION;
+      }
+      VERSION.set(detectedVersion);
+    }
+
+    return VERSION.get();
+  }
+
+  @VisibleForTesting
+  static String versionFromJar() {
+    // Choose {@link DataStream} class because it is one of the core Flink API
+    return DataStream.class.getPackage().getImplementationVersion();
+  }
+
+  @VisibleForTesting
+  static void setVersion(String version) {
+    VERSION.set(version);
   }
 }
diff --git a/flink/v1.17/flink/src/test/java/org/apache/iceberg/flink/util/TestFlinkPackage.java b/flink/v1.17/flink/src/test/java/org/apache/iceberg/flink/util/TestFlinkPackage.java
@@ -20,6 +20,8 @@
 
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
 public class TestFlinkPackage {
 
@@ -28,4 +30,25 @@ public class TestFlinkPackage {
   public void testVersion() {
     Assert.assertEquals("1.17.1", FlinkPackage.version());
   }
+
+  @Test
+  public void testDefaultVersion() {
+    // It's difficult to reproduce a reflection error in a unit test, so we just inject a mocked
+    // fault to test the default logic
+
+    // First make sure we're not caching a version result from a previous test
+    FlinkPackage.setVersion(null);
+    try (MockedStatic<FlinkPackage> mockedStatic = Mockito.mockStatic(FlinkPackage.class)) {
+      mockedStatic.when(FlinkPackage::versionFromJar).thenThrow(RuntimeException.class);
+      mockedStatic.when(FlinkPackage::version).thenCallRealMethod();
+      Assert.assertEquals(FlinkPackage.FLINK_UNKNOWN_VERSION, FlinkPackage.version());
+    }
+    FlinkPackage.setVersion(null);
+    try (MockedStatic<FlinkPackage> mockedStatic = Mockito.mockStatic(FlinkPackage.class)) {
+      mockedStatic.when(FlinkPackage::versionFromJar).thenReturn(null);
+      mockedStatic.when(FlinkPackage::version).thenCallRealMethod();
+      FlinkPackage.setVersion(null);
+      Assert.assertEquals(FlinkPackage.FLINK_UNKNOWN_VERSION, FlinkPackage.version());
+    }
+  }
 }
diff --git a/flink/v1.18/flink/src/main/java/org/apache/iceberg/flink/util/FlinkPackage.java b/flink/v1.18/flink/src/main/java/org/apache/iceberg/flink/util/FlinkPackage.java
@@ -18,16 +18,44 @@
  */
 package org.apache.iceberg.flink.util;
 
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
 
 public class FlinkPackage {
-  /** Choose {@link DataStream} class because it is one of the core Flink API. */
-  private static final String VERSION = DataStream.class.getPackage().getImplementationVersion();
+
+  private static final AtomicReference<String> VERSION = new AtomicReference<>();
+  public static final String FLINK_UNKNOWN_VERSION = "FLINK-UNKNOWN-VERSION";
 
   private FlinkPackage() {}
 
   /** Returns Flink version string like x.y.z */
   public static String version() {
-    return VERSION;
+    if (null == VERSION.get()) {
+      String detectedVersion = null;
+      try {
+        detectedVersion = versionFromJar();
+        // use unknown version in case exact implementation version can't be found from the jar
+        // (this can happen if the DataStream class appears multiple times in the same classpath
+        // such as with shading)
+        detectedVersion = detectedVersion != null ? detectedVersion : FLINK_UNKNOWN_VERSION;
+      } catch (Exception e) {
+        detectedVersion = FLINK_UNKNOWN_VERSION;
+      }
+      VERSION.set(detectedVersion);
+    }
+
+    return VERSION.get();
+  }
+
+  @VisibleForTesting
+  static String versionFromJar() {
+    // Choose {@link DataStream} class because it is one of the core Flink API
+    return DataStream.class.getPackage().getImplementationVersion();
+  }
+
+  @VisibleForTesting
+  static void setVersion(String version) {
+    VERSION.set(version);
   }
 }
diff --git a/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/util/TestFlinkPackage.java b/flink/v1.18/flink/src/test/java/org/apache/iceberg/flink/util/TestFlinkPackage.java
@@ -20,6 +20,8 @@
 
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
 public class TestFlinkPackage {
 
@@ -28,4 +30,25 @@ public class TestFlinkPackage {
   public void testVersion() {
     Assert.assertEquals("1.18.1", FlinkPackage.version());
   }
+
+  @Test
+  public void testDefaultVersion() {
+    // It's difficult to reproduce a reflection error in a unit test, so we just inject a mocked
+    // fault to test the default logic
+
+    // First make sure we're not caching a version result from a previous test
+    FlinkPackage.setVersion(null);
+    try (MockedStatic<FlinkPackage> mockedStatic = Mockito.mockStatic(FlinkPackage.class)) {
+      mockedStatic.when(FlinkPackage::versionFromJar).thenThrow(RuntimeException.class);
+      mockedStatic.when(FlinkPackage::version).thenCallRealMethod();
+      Assert.assertEquals(FlinkPackage.FLINK_UNKNOWN_VERSION, FlinkPackage.version());
+    }
+    FlinkPackage.setVersion(null);
+    try (MockedStatic<FlinkPackage> mockedStatic = Mockito.mockStatic(FlinkPackage.class)) {
+      mockedStatic.when(FlinkPackage::versionFromJar).thenReturn(null);
+      mockedStatic.when(FlinkPackage::version).thenCallRealMethod();
+      FlinkPackage.setVersion(null);
+      Assert.assertEquals(FlinkPackage.FLINK_UNKNOWN_VERSION, FlinkPackage.version());
+    }
+  }
 }