Add support for extensions to depend upon other extensions

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/annotations/ExtensionDependsOn.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.annotations;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Documented
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE})
+public @interface ExtensionDependsOn {
+    /**
+     * The list of classes that this extension depends on
+     * @return Array of Class objects representing the classes this extension depends on
+     */
+    Class<?>[] dependentClasses() default {};
+}

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/annotations/ExtensionProvides.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.annotations;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation to specify that a class provides an extension and its dependencies.
+ * This annotation can be used to declare what extension points a class provides
+ * and what other extension points it depends on.
+ */
+@Documented
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE})
+public @interface ExtensionProvides {
+    /**
+     * The list of classes that this extension provides
+     * @return Array of Class objects representing the classes this extension provides
+     */
+    Class<?>[] providedClasses() default {};
+}

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/plugin/ExtensionPoints.java

@@ -20,4 +20,6 @@ public interface ExtensionPoints {
      * @since 2.3
      */
     void addExtensionProvider(ExtensionProvider<?> extensionProvider);
+
+    <T> T getExtensionProvider(Class<T> type);
 }

data-prepper-plugin-framework/src/main/java/org/opensearch/dataprepper/plugin/DataPrepperExtensionPoints.java

@@ -50,6 +50,12 @@ public void addExtensionProvider(final ExtensionProvider extensionProvider) {
         providerClassesSet.add(extensionProvider.supportedClass());
     }
 
+    @Override
+    public <T> T getExtensionProvider(final Class<T> type) {
+        sharedApplicationContext.refresh();
+        return sharedApplicationContext.getBean(type);
+    }
+
     private static class EmptyContext implements ExtensionProvider.Context {
 
     }

data-prepper-plugin-framework/src/main/java/org/opensearch/dataprepper/plugin/ExtensionLoader.java

@@ -6,6 +6,8 @@
 package org.opensearch.dataprepper.plugin;
 
 import org.opensearch.dataprepper.model.annotations.DataPrepperExtensionPlugin;
+import org.opensearch.dataprepper.model.annotations.ExtensionDependsOn;
+import org.opensearch.dataprepper.model.annotations.ExtensionProvides;
 import org.opensearch.dataprepper.model.configuration.PipelinesDataFlowModel;
 import org.opensearch.dataprepper.model.plugin.ExtensionPlugin;
 import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException;
@@ -27,10 +29,17 @@ public class ExtensionLoader {
     public class ExtensionPluginWithContext {
         ExtensionPlugin extensionPlugin;
         boolean configured;
+        Class<?>[] dependentClasses;
+        Class<?>[] providedClasses;
 
-        public ExtensionPluginWithContext(final ExtensionPlugin extensionPlugin, final boolean isConfigured) {
+        public ExtensionPluginWithContext(final ExtensionPlugin extensionPlugin,
+                                          final boolean isConfigured,
+                                          final Class<?>[] dependentClasses,
+                                          final Class<?>[] providedClasses) {
             this.extensionPlugin = extensionPlugin;
             this.configured = isConfigured;
+            this.dependentClasses = dependentClasses;
+            this.providedClasses = providedClasses;
         }
 
         public ExtensionPlugin getExtensionPlugin() {
@@ -40,6 +49,14 @@ public ExtensionPlugin getExtensionPlugin() {
         public boolean isConfigured() {
             return configured;
         }
+
+        public Class<?>[] getDependentClasses() {
+            return dependentClasses;
+        }
+
+        public Class<?>[] getProvidedClasses() {
+            return providedClasses;
+        }
     }
 
     private Comparator<ExtensionLoader.ExtensionPluginWithContext> extensionsLoaderComparator;
@@ -73,9 +90,17 @@ public List<? extends ExtensionPlugin> loadExtensions() {
                     final String pluginName = convertClassToName(extensionClass);
                     try {
                         final PluginArgumentsContext pluginArgumentsContext = getConstructionContext(extensionClass);
+                        final ExtensionProvides extensionProvidesAnnotation = extensionClass.getAnnotation(ExtensionProvides.class);
+                        final ExtensionDependsOn extensionDependsOnAnnotation = extensionClass.getAnnotation(ExtensionDependsOn.class);
+
+                        final Class<?>[] providedClasses = extensionProvidesAnnotation != null ?
+                                extensionProvidesAnnotation.providedClasses() : new Class<?>[]{};
+                        final Class<?>[] dependentClasses = extensionDependsOnAnnotation != null ?
+                                extensionDependsOnAnnotation.dependentClasses() : new Class<?>[]{};
+
                         final Object config = pluginArgumentsContext.getArgument(0);
                         return new ExtensionPluginWithContext(extensionPluginCreator.newPluginInstance(
-                                extensionClass, pluginArgumentsContext, pluginName), (config != null));
+                                extensionClass, pluginArgumentsContext, pluginName), (config != null), dependentClasses, providedClasses);
                     } catch (Exception e) {
                         final PluginError pluginError = PluginError.builder()
                                 .componentType(PipelinesDataFlowModel.EXTENSION_PLUGIN_TYPE)

data-prepper-plugin-framework/src/main/java/org/opensearch/dataprepper/plugin/PluginCreatorContext.java

@@ -4,6 +4,7 @@
 
 import javax.inject.Named;
 
+import java.util.Arrays;
 import java.util.Comparator;
 
 @Named
@@ -21,6 +22,35 @@ public PluginCreator pluginCreator(
 
     @Bean(name = "extensionsLoaderComparator")
     public Comparator<ExtensionLoader.ExtensionPluginWithContext> extensionsLoaderComparator() {
-        return Comparator.comparing(ExtensionLoader.ExtensionPluginWithContext::isConfigured).reversed();
+        return (extensionOne, extensionTwo) -> {
+            // First, compare by configuration status (configured ones first)
+            int configCompare = Boolean.compare(extensionTwo.isConfigured(), extensionOne.isConfigured());
+            if (configCompare != 0) {
+                return configCompare;
+            }
+
+            // Get the provided and dependent classes for both extensions
+            Class<?>[] extensionOneProvidedClasses = extensionOne.getProvidedClasses();
+            Class<?>[] extensionTwoProvidedClasses = extensionTwo.getProvidedClasses();
+            Class<?>[] extensionOneDependentClasses = extensionOne.getDependentClasses();
+            Class<?>[] extensionTwoDependentClasses = extensionTwo.getDependentClasses();
+
+            // If extensionOne provides any classes that extensionTwo depends on, extensionOne should go first
+            if (containsAnyExtensionDependencies(extensionOneProvidedClasses, extensionTwoDependentClasses)) {
+                return -1;
+            }
+
+            // If extensionTwo provides any classes that extensionOne depends on, extensionTwo should go first
+            if (containsAnyExtensionDependencies(extensionTwoProvidedClasses, extensionOneDependentClasses)) {
+                return 1;
+            }
+
+            return 0;
+        };
+    }
+
+    private boolean containsAnyExtensionDependencies(final Class<?>[] provided, final Class<?>[] dependencies) {
+        return Arrays.stream(dependencies).anyMatch(dep ->
+                Arrays.asList(provided).contains(dep));
     }
 }

data-prepper-plugin-framework/src/test/java/org/opensearch/dataprepper/plugin/DataPrepperExtensionPointsTest.java

@@ -127,6 +127,20 @@ void addExtensionProvider_should_registerBean_as_prototype() {
         verifyRegisterBeanAsPrototype(coreApplicationContext);
     }
 
+    @Test
+    void getExtensionProvider_refreshes_shared_context_and_returns_correct_bean() {
+        final Class<DefaultPluginFactory> defaultPluginFactoryClass = DefaultPluginFactory.class;
+        final DefaultPluginFactory defaultPluginFactory = mock(DefaultPluginFactory.class);
+
+        when(sharedApplicationContext.getBean(defaultPluginFactoryClass)).thenReturn(defaultPluginFactory);
+
+        final DefaultPluginFactory result = createObjectUnderTest().getExtensionProvider(defaultPluginFactoryClass);
+
+        assertThat(result, equalTo(defaultPluginFactory));
+
+        verify(sharedApplicationContext).refresh();
+    }
+
     private void verifyRegisterBeanWithProvideInstance(final GenericApplicationContext applicationContext) {
         reset(extensionProvider);
         final ArgumentCaptor<Supplier<Object>> supplierArgumentCaptor =

data-prepper-plugin-framework/src/test/java/org/opensearch/dataprepper/plugin/ExtensionLoaderTest.java

@@ -18,6 +18,8 @@
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.dataprepper.model.annotations.ExtensionDependsOn;
+import org.opensearch.dataprepper.model.annotations.ExtensionProvides;
 import org.opensearch.dataprepper.model.configuration.PipelinesDataFlowModel;
 import org.opensearch.dataprepper.model.plugin.ExtensionPlugin;
 import org.opensearch.dataprepper.model.plugin.InvalidPluginConfigurationException;
@@ -45,6 +47,7 @@
 import static org.hamcrest.CoreMatchers.nullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -224,7 +227,7 @@ void loadExtensions_throws_InvalidPluginConfigurationException_when_extensionPlu
     }
 
     @Test
-    void loadExtensions_returns_multiple_extensions_for_multiple_plugin_classes() {
+    void loadExtensions_returns_multiple_extensions_for_multiple_plugin_classes_in_correct_order() {
         final Collection<Class<? extends ExtensionPlugin>> pluginClasses = new HashSet<>();
         final Collection<ExtensionPlugin> expectedPlugins = new ArrayList<>();
 
@@ -256,6 +259,13 @@ void loadExtensions_returns_multiple_extensions_for_multiple_plugin_classes() {
         for (ExtensionPlugin expectedPlugin : actualPlugins) {
             assertThat(actualPlugins, hasItem(expectedPlugin));
         }
+
+        assertThat(actualPlugins.get(0), instanceOf(TestExtension2.class));
+        assertTrue(actualPlugins.get(1) instanceof TestExtension1 || actualPlugins.get(1) instanceof TestExtension3,
+                "Expected result to be either TestExtension1 or TestExtension3 but was " + actualPlugins.get(1).getClass().getName());
+
+        assertTrue(actualPlugins.get(2) instanceof TestExtension1 || actualPlugins.get(2) instanceof TestExtension3,
+                "Expected result to be either TestExtension1 or TestExtension3 but was " + actualPlugins.get(2).getClass().getName());
         assertThat(pluginErrorCollector.getPluginErrors().isEmpty(), is(true));
     }
 
@@ -335,10 +345,15 @@ private static Stream<Arguments>  validExtensionConfigs() {
                 null);
     }
 
+    @ExtensionDependsOn(dependentClasses = TestExtensionConfig.class)
     private interface TestExtension1 extends ExtensionPlugin {
     }
+
+    @ExtensionProvides(providedClasses = TestExtensionConfig.class)
     private interface TestExtension2 extends ExtensionPlugin {
     }
+
+    @ExtensionDependsOn(dependentClasses = TestExtensionConfig.class)
     private interface TestExtension3 extends ExtensionPlugin {
     }
 

data-prepper-plugin-framework/src/test/java/org/opensearch/dataprepper/plugin/PluginCreatorContextTest.java

@@ -40,16 +40,29 @@ public void test_pluginCreator() {
 
     @Test
     public void test_extensionsLoaderComparator() {
+        final Class<?>[] classes = {DefaultPluginFactory.class};
+
         ExtensionLoader.ExtensionPluginWithContext context1 = mock(ExtensionLoader.ExtensionPluginWithContext.class);
         ExtensionLoader.ExtensionPluginWithContext context2 = mock(ExtensionLoader.ExtensionPluginWithContext.class);
         Comparator<ExtensionLoader.ExtensionPluginWithContext> extensionsLoaderComparator = pluginCreatorContext.extensionsLoaderComparator();
         assertNotNull(extensionsLoaderComparator);
         when(context1.isConfigured()).thenReturn(true);
+        when(context1.getDependentClasses()).thenReturn(classes);
+        when(context1.getProvidedClasses()).thenReturn(new Class<?>[]{});
+
         when(context2.isConfigured()).thenReturn(true);
-        assertThat(extensionsLoaderComparator.compare(context1, context2), equalTo(0));
+        when(context2.getProvidedClasses()).thenReturn(classes);
+        when(context2.getDependentClasses()).thenReturn(new Class<?>[]{});
+        assertThat(extensionsLoaderComparator.compare(context1, context2), equalTo(1));
+
         when(context1.isConfigured()).thenReturn(false);
+        when(context1.getDependentClasses()).thenReturn(new Class<?>[]{});
+        when(context1.getProvidedClasses()).thenReturn(classes);
+
         when(context2.isConfigured()).thenReturn(false);
-        assertThat(extensionsLoaderComparator.compare(context1, context2), equalTo(0));
+        when(context2.getProvidedClasses()).thenReturn(new Class<?>[]{});
+        when(context2.getDependentClasses()).thenReturn(classes);
+        assertThat(extensionsLoaderComparator.compare(context1, context2), equalTo(-1));
         when(context1.isConfigured()).thenReturn(false);
         when(context2.isConfigured()).thenReturn(true);
         assertThat(extensionsLoaderComparator.compare(context1, context2), greaterThan(0));

data-prepper-plugins/aws-plugin/src/main/java/org/opensearch/dataprepper/plugins/aws/AwsPlugin.java

@@ -5,8 +5,10 @@
 
 package org.opensearch.dataprepper.plugins.aws;
 
+import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
 import org.opensearch.dataprepper.model.annotations.DataPrepperExtensionPlugin;
 import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor;
+import org.opensearch.dataprepper.model.annotations.ExtensionProvides;
 import org.opensearch.dataprepper.model.plugin.ExtensionPlugin;
 import org.opensearch.dataprepper.model.plugin.ExtensionPoints;
 
@@ -15,6 +17,7 @@
  * Data Prepper as an extension plugin. Everything starts from here.
  */
 @DataPrepperExtensionPlugin(modelType = AwsPluginConfig.class, rootKeyJsonPath = "/aws/configurations")
+@ExtensionProvides(providedClasses = {AwsCredentialsSupplier.class})
 public class AwsPlugin implements ExtensionPlugin {
     private final DefaultAwsCredentialsSupplier defaultAwsCredentialsSupplier;