From 4074ec37858b5fc42074468c981d507cdae4d402 Mon Sep 17 00:00:00 2001
From: danbress <dbress@onyxconsults.com>
Date: Sat, 30 May 2015 12:04:46 -0400
Subject: [PATCH] NIFI-632 using NarThreadContextClassLoader when generating
 documentation

---
 .../nifi-framework/nifi-documentation/pom.xml |  4 +++
 .../nifi/documentation/DocGenerator.java      | 27 ++++++++++++-------
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/pom.xml b/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/pom.xml
index 340f2c54b3..66febbbb73 100644
--- a/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/pom.xml
+++ b/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/pom.xml
@@ -26,6 +26,10 @@
             <groupId>org.apache.nifi</groupId>
             <artifactId>nifi-api</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.apache.nifi</groupId>
+            <artifactId>nifi-framework-core-api</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.apache.nifi</groupId>
             <artifactId>nifi-properties</artifactId>
diff --git a/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/src/main/java/org/apache/nifi/documentation/DocGenerator.java b/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/src/main/java/org/apache/nifi/documentation/DocGenerator.java
index 0e9ed37414..c817bfa11a 100644
--- a/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/src/main/java/org/apache/nifi/documentation/DocGenerator.java
+++ b/nifi/nifi-nar-bundles/nifi-framework-bundle/nifi-framework/nifi-documentation/src/main/java/org/apache/nifi/documentation/DocGenerator.java
@@ -33,6 +33,7 @@ import org.apache.nifi.documentation.init.ControllerServiceInitializer;
 import org.apache.nifi.documentation.init.ProcessorInitializer;
 import org.apache.nifi.documentation.init.ReportingTaskingInitializer;
 import org.apache.nifi.nar.ExtensionManager;
+import org.apache.nifi.nar.NarThreadContextClassLoader;
 import org.apache.nifi.processor.Processor;
 import org.apache.nifi.reporting.InitializationException;
 import org.apache.nifi.reporting.ReportingTask;
@@ -68,17 +69,25 @@ public class DocGenerator {
         logger.debug("Generating documentation for: " + extensionClasses.size() + " components in: "
                 + explodedNiFiDocsDir);
 
-        for (final Class<?> extensionClass : extensionClasses) {
-            if (ConfigurableComponent.class.isAssignableFrom(extensionClass)) {
-                final Class<? extends ConfigurableComponent> componentClass = extensionClass
-                        .asSubclass(ConfigurableComponent.class);
-                try {
-                    logger.debug("Documenting: " + componentClass);
-                    document(explodedNiFiDocsDir, componentClass);
-                } catch (Exception e) {
-                    logger.warn("Unable to document: " + componentClass, e);
+        // save the original class loader and use a NarThreadContextClassLoader
+        final ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();
+        try {
+            Thread.currentThread().setContextClassLoader(NarThreadContextClassLoader.getInstance());
+
+            for (final Class<?> extensionClass : extensionClasses) {
+                if (ConfigurableComponent.class.isAssignableFrom(extensionClass)) {
+                    final Class<? extends ConfigurableComponent> componentClass = extensionClass
+                            .asSubclass(ConfigurableComponent.class);
+                    try {
+                        logger.debug("Documenting: " + componentClass);
+                        document(explodedNiFiDocsDir, componentClass);
+                    } catch (Exception e) {
+                        logger.warn("Unable to document: " + componentClass, e);
+                    }
                 }
             }
+        } finally {
+            Thread.currentThread().setContextClassLoader(originalClassLoader);
         }
     }