Set avro class-loader on the current thread

Velcity engine loads classes through the thread class loader
sbt · Jan 14, 2025 · bd2011b · bd2011b
1 parent ed72061
commit bd2011b
Showing 1 changed file with 24 additions and 21 deletions.
diff --git a/plugin/src/main/scala/com/github/sbt/avro/SbtAvro.scala b/plugin/src/main/scala/com/github/sbt/avro/SbtAvro.scala
@@ -251,39 +251,42 @@ object SbtAvro extends AutoPlugin {
                     .toArray,
                   this.getClass.getClassLoader
                 )
+                val initLoader = Thread.currentThread().getContextClassLoader
 
-                val compiler = avroClassLoader
-                  .loadClass(avroCompiler.value)
-                  .getDeclaredConstructor()
-                  .newInstance()
-                  .asInstanceOf[AvroCompiler]
-
-                compiler.setStringType(avroStringType.value)
-                compiler.setFieldVisibility(avroFieldVisibility.value.toUpperCase)
-                compiler.setEnableDecimalLogicalType(avroEnableDecimalLogicalType.value)
-                compiler.setCreateSetters(avroCreateSetters.value)
-                compiler.setOptionalGetters(avroOptionalGetters.value)
-
-                val recs = records.map(avroClassLoader.loadClass)
-                val avdls = srcDirs.flatMap(d => (d ** AvroAvdlFilter).get())
-                val avscs = srcDirs.flatMap(d => (d ** AvroAvscFilter).get())
-                val avprs = srcDirs.flatMap(d => (d ** AvroAvrpFilter).get())
-
-                out.log.info(
-                  s"Avro compiler ${avroVersion.value} using stringType=${avroStringType.value}"
-                )
                 try {
+                  val compiler = avroClassLoader
+                    .loadClass(avroCompiler.value)
+                    .getDeclaredConstructor()
+                    .newInstance()
+                    .asInstanceOf[AvroCompiler]
+
+                  compiler.setStringType(avroStringType.value)
+                  compiler.setFieldVisibility(avroFieldVisibility.value.toUpperCase)
+                  compiler.setEnableDecimalLogicalType(avroEnableDecimalLogicalType.value)
+                  compiler.setCreateSetters(avroCreateSetters.value)
+                  compiler.setOptionalGetters(avroOptionalGetters.value)
+
+                  val recs = records.map(avroClassLoader.loadClass)
+                  val avdls = srcDirs.flatMap(d => (d ** AvroAvdlFilter).get())
+                  val avscs = srcDirs.flatMap(d => (d ** AvroAvscFilter).get())
+                  val avprs = srcDirs.flatMap(d => (d ** AvroAvrpFilter).get())
+
+                  out.log.info(
+                    s"Avro compiler ${avroVersion.value} using stringType=${avroStringType.value}"
+                  )
+                  Thread.currentThread().setContextClassLoader(avroClassLoader)
                   compiler.recompile(recs.toArray, outDir)
                   compiler.compileAvscs(avscs.toArray, outDir)
                   compiler.compileIdls(avdls.toArray, outDir)
                   compiler.compileAvprs(avprs.toArray, outDir)
 
                   (outDir ** SbtAvro.JavaFileFilter).get().toSet
                 } catch {
-                  case e: Exception =>
+                  case e: RuntimeException =>
                     out.log.err(e.getMessage)
                     throw new AvroGenerateFailedException
                 } finally {
+                  Thread.currentThread().setContextClassLoader(initLoader)
                   avroClassLoader.close()
                 }
               } else {