change to match parameter type length

AbsaOSS · Sep 2, 2024 · 4925c68 · 4925c68
1 parent 36c4f7b
commit 4925c68
Showing 1 changed file with 13 additions and 33 deletions.
diff --git a/core/src/main/scala/za/co/absa/spline/harvester/plugin/embedded/BigQueryPlugin.scala b/core/src/main/scala/za/co/absa/spline/harvester/plugin/embedded/BigQueryPlugin.scala
@@ -16,12 +16,7 @@
 
 package za.co.absa.spline.harvester.plugin.embedded
 
-import java.lang.reflect.Method
-import java.lang.reflect.Modifier.isStatic
-import java.util.Optional
-
 import io.github.classgraph.ClassGraph
-import javax.annotation.Priority
 import org.apache.hadoop.conf.Configuration
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.execution.datasources.{LogicalRelation, SaveIntoDataSourceCommand}
@@ -36,6 +31,10 @@ import za.co.absa.spline.harvester.plugin.embedded.BigQueryPlugin.SparkBigQueryC
 import za.co.absa.spline.harvester.plugin.embedded.BigQueryPlugin._
 import za.co.absa.spline.harvester.plugin.{BaseRelationProcessing, Plugin, RelationProviderProcessing}
 
+import java.lang.reflect.Method
+import java.lang.reflect.Modifier.isStatic
+import java.util.Optional
+import javax.annotation.Priority
 import scala.collection.JavaConverters._
 import scala.language.reflectiveCalls
 
@@ -143,33 +142,14 @@ object BigQueryPlugin {
       def getParentProjectId: String
     }
     private val clazz = findPossiblyShadedClass("com.google.cloud", "com.google.cloud.spark.bigquery.SparkBigQueryConfig")
-    private val methodFrom7: Option[Method] = clazz
-      .getMethods
-      .find(
-        m => m.getName == "from"
-          && isStatic(m.getModifiers)
-          && m.getParameterTypes.length == 7
-          && m.getReturnType.getSimpleName == "SparkBigQueryConfig"
-      )
-    private val methodFrom8: Option[Method] = clazz
+    private val methodFrom: Method = clazz
       .getMethods
       .find(
         m => m.getName == "from"
           && isStatic(m.getModifiers)
-          && m.getParameterTypes.length == 8
+          && (m.getParameterTypes.length == 7 || m.getParameterTypes.length == 8 || m.getParameterTypes.length == 9)
           && m.getReturnType.getSimpleName == "SparkBigQueryConfig"
       )
-    private val methodFrom9: Option[Method] = clazz
-      .getMethods
-      .find(
-        m => m.getName == "from"
-          && isStatic(m.getModifiers)
-          && m.getParameterTypes.length == 9
-          && m.getReturnType.getSimpleName == "SparkBigQueryConfig"
-      )
-    private val methodFrom: Method = methodFrom7
-      .orElse(methodFrom8)
-      .orElse(methodFrom9)
       .getOrElse(sys.error(s"Cannot find method `public static SparkBigQueryConfig from(... {7|8|9 args} ...)` in the class `$clazz`"))
 
     object ImmutableMap {
@@ -188,13 +168,13 @@ object BigQueryPlugin {
     }
 
     val from: (java.util.Map[_, _], ImmutableMap.ImmutableMap, Configuration, ImmutableMap.ImmutableMap, Integer, SQLConf, String, Optional[StructType], java.lang.Boolean) => SparkBigQueryConfig = {
-      case (a, b, c, d, e, f, g, h, i) =>
-        if (methodFrom7.isDefined)
-          methodFrom7.get.invoke(clazz, a, b, c, e, f, g, h).asInstanceOf[SparkBigQueryConfig]
-        else if (methodFrom8.isDefined)
-          methodFrom8.get.invoke(clazz, a, b, c, e, f, g, h, i).asInstanceOf[SparkBigQueryConfig]
-        else
-          methodFrom9.get.invoke(clazz, a, b, c, d, e, f, g, h, i).asInstanceOf[SparkBigQueryConfig]
+      case (a, b, c, d, e, f, g, h, i) => {
+        methodFrom.getParameterTypes.length match {
+          case 7 => methodFrom.invoke(clazz, a, b, c, e, f, g, h).asInstanceOf[SparkBigQueryConfig]
+          case 8 => methodFrom.invoke(clazz, a, b, c, e, f, g, h, i).asInstanceOf[SparkBigQueryConfig]
+          case 9 => methodFrom.invoke(clazz, a, b, c, d, e, f, g, h, i).asInstanceOf[SparkBigQueryConfig]
+        }
+      }
     }
   }