ENH: introduce primary fields in plugin schemas

data-prepper-plugin-schema-cli/README.md

@@ -5,9 +5,21 @@ This module includes the SDK and CLI for generating schemas for Data Prepper pip
 ## CLI Usage
 
 ```
-./gradlew :data-prepper-plugin-schema-cli:run --args='--plugin_type=processor --plugin_names=grok --output_folder=/path/to/schemas'
+./gradlew :data-prepper-plugin-schema-cli:run --args='--plugin_type=processor --plugin_names=grok --output_folder=/path/to/schemas --primary_fields_override /path/to/override.yaml'
 ```
 
 * plugin_type: A required parameter specifies type of processor. Valid options are `source`, `buffer`, `processor`, `route`, `sink`.
 * plugin_names: An optional parameter filters the result by plugin names separated by `,`, e.g. `grok,date`.
 * output_folder: An optional parameter to specify the output folder path.
+* primary_fields_override: An optional parameter to specify the custom JSON/YAML file path which includes primary fields override for specified plugin names. e.g.
+
+primary_fields_override.yaml
+```
+key_value: [ "source" ]
+```
+will generate override result on primary_fields in `key_value` processor:
+```
+{
+    "primary_fields": [ "source" ]
+}
+```

data-prepper-plugin-schema-cli/build.gradle

@@ -12,6 +12,7 @@ dependencies {
     implementation project(':data-prepper-plugin-framework')
     implementation project(':data-prepper-plugin-schema')
     implementation 'com.fasterxml.jackson.core:jackson-databind'
+    implementation 'com.fasterxml.jackson.dataformat:jackson-dataformat-yaml'
     implementation 'org.reflections:reflections:0.10.2'
     implementation 'com.github.victools:jsonschema-maven-plugin:4.35.0'
     implementation 'com.github.victools:jsonschema-generator:4.35.0'

data-prepper-plugin-schema-cli/src/main/java/org/opensearch/dataprepper/schemas/DataPrepperPluginSchemaExecute.java

@@ -1,5 +1,7 @@
 package org.opensearch.dataprepper.schemas;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
 import com.github.victools.jsonschema.generator.OptionPreset;
 import com.github.victools.jsonschema.generator.SchemaVersion;
 import org.opensearch.dataprepper.plugin.ClasspathPluginProvider;
@@ -9,7 +11,10 @@
 import org.slf4j.LoggerFactory;
 import picocli.CommandLine;
 
+import java.io.File;
 import java.io.IOException;
+import java.net.URISyntaxException;
+import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -21,6 +26,7 @@
 
 public class DataPrepperPluginSchemaExecute implements Runnable {
     private static final Logger LOG = LoggerFactory.getLogger(DataPrepperPluginSchemaExecute.class);
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(new YAMLFactory());
     static final String DEFAULT_PLUGINS_CLASSPATH = "org.opensearch.dataprepper.plugins";
 
     @CommandLine.Option(names = {"--plugin_type"}, required = true)
@@ -29,6 +35,9 @@ public class DataPrepperPluginSchemaExecute implements Runnable {
     @CommandLine.Option(names = {"--plugin_names"})
     private String pluginNames;
 
+    @CommandLine.Option(names = {"--primary_fields_override"})
+    private String primaryFieldsOverrideFilePath;
+
     @CommandLine.Option(names = {"--site.url"}, defaultValue = "https://opensearch.org")
     private String siteUrl;
     @CommandLine.Option(names = {"--site.baseurl"}, defaultValue = "/docs/latest")
@@ -45,8 +54,16 @@ public static void main(String[] args) {
     @Override
     public void run() {
         final PluginProvider pluginProvider = new ClasspathPluginProvider();
+        final PrimaryFieldsOverride primaryFieldsOverride;
+        try {
+            primaryFieldsOverride = primaryFieldsOverrideFilePath == null ? new PrimaryFieldsOverride() :
+                     OBJECT_MAPPER.readValue(new File(primaryFieldsOverrideFilePath), PrimaryFieldsOverride.class);
+        } catch (IOException e) {
+            throw new RuntimeException("primary fields override filepath does not exist. ", e);
+        }
         final PluginConfigsJsonSchemaConverter pluginConfigsJsonSchemaConverter = new PluginConfigsJsonSchemaConverter(
-                pluginProvider, new JsonSchemaConverter(DataPrepperModules.dataPrepperModules(), pluginProvider), siteUrl, siteBaseUrl);
+                pluginProvider, new JsonSchemaConverter(DataPrepperModules.dataPrepperModules(), pluginProvider),
+                primaryFieldsOverride, siteUrl, siteBaseUrl);
         final Class<?> pluginType = pluginConfigsJsonSchemaConverter.pluginTypeNameToPluginType(pluginTypeName);
         final Map<String, String> pluginNameToJsonSchemaMap = pluginConfigsJsonSchemaConverter.convertPluginConfigsIntoJsonSchemas(
                 SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON, pluginType);

data-prepper-plugin-schema/src/main/java/org/opensearch/dataprepper/schemas/PluginConfigsJsonSchemaConverter.java

@@ -1,5 +1,6 @@
 package org.opensearch.dataprepper.schemas;
 
+import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.github.victools.jsonschema.generator.OptionPreset;
 import com.github.victools.jsonschema.generator.SchemaVersion;
@@ -31,6 +32,7 @@ public class PluginConfigsJsonSchemaConverter {
     static final String SITE_BASE_URL_PLACEHOLDER = "{{site.baseurl}}";
     static final String DOCUMENTATION_LINK_KEY = "documentation";
     static final String PLUGIN_NAME_KEY = "name";
+    static final String PRIMARY_FIELDS_KEY = "primary_fields";
     static final String PLUGIN_DOCUMENTATION_URL_FORMAT =
             "%s%s/data-prepper/pipelines/configuration/%s/%s/";
     static final Map<Class<?>, String> PLUGIN_TYPE_TO_URI_PARAMETER_MAP = Map.of(
@@ -52,14 +54,17 @@ public class PluginConfigsJsonSchemaConverter {
     private final String siteBaseUrl;
     private final PluginProvider pluginProvider;
     private final JsonSchemaConverter jsonSchemaConverter;
+    private final PrimaryFieldsOverride primaryFieldsOverride;
 
     public PluginConfigsJsonSchemaConverter(
             final PluginProvider pluginProvider,
             final JsonSchemaConverter jsonSchemaConverter,
+            final PrimaryFieldsOverride primaryFieldsOverride,
             final String siteUrl,
             final String siteBaseUrl) {
         this.pluginProvider = pluginProvider;
         this.jsonSchemaConverter = jsonSchemaConverter;
+        this.primaryFieldsOverride = primaryFieldsOverride;
         this.siteUrl = siteUrl == null ? SITE_URL_PLACEHOLDER : siteUrl;
         this.siteBaseUrl = siteBaseUrl == null ? SITE_BASE_URL_PLACEHOLDER : siteBaseUrl;
     }
@@ -87,6 +92,7 @@ public Map<String, String> convertPluginConfigsIntoJsonSchemas(
                         final ObjectNode jsonSchemaNode = jsonSchemaConverter.convertIntoJsonSchema(
                                 schemaVersion, optionPreset, entry.getValue());
                         addPluginName(jsonSchemaNode, pluginName);
+                        addPrimaryFields(jsonSchemaNode, pluginName);
                         addDocumentationLink(jsonSchemaNode, pluginName, pluginType);
                         value = jsonSchemaNode.toPrettyString();
                     } catch (final Exception e) {
@@ -131,4 +137,10 @@ private void addPluginName(final ObjectNode jsonSchemaNode,
                                final String pluginName) {
         jsonSchemaNode.put(PLUGIN_NAME_KEY, pluginName);
     }
+
+    private void addPrimaryFields(final ObjectNode jsonSchemaNode,
+                                  final String pluginName) {
+        final ArrayNode primaryFieldsNode = jsonSchemaNode.putArray(PRIMARY_FIELDS_KEY);
+        primaryFieldsOverride.getPrimaryFieldsForComponent(pluginName).forEach(primaryFieldsNode::add);
+    }
 }

data-prepper-plugin-schema/src/main/java/org/opensearch/dataprepper/schemas/PrimaryFieldsOverride.java

@@ -0,0 +1,24 @@
+package org.opensearch.dataprepper.schemas;
+
+import com.fasterxml.jackson.annotation.JsonAnySetter;
+import com.fasterxml.jackson.annotation.JsonCreator;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+public class PrimaryFieldsOverride {
+
+    @JsonAnySetter
+    private final Map<String, Set<String>> primaryFieldsMap;
+
+    @JsonCreator
+    public PrimaryFieldsOverride() {
+        primaryFieldsMap = new HashMap<>();
+    }
+
+    public Set<String> getPrimaryFieldsForComponent(final String componentName) {
+        return primaryFieldsMap.getOrDefault(componentName, Collections.emptySet());
+    }
+}

data-prepper-plugin-schema/src/test/java/org/opensearch/dataprepper/schemas/PluginConfigsJsonSchemaConverterIT.java

@@ -27,6 +27,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.opensearch.dataprepper.schemas.PluginConfigsJsonSchemaConverter.DOCUMENTATION_LINK_KEY;
 import static org.opensearch.dataprepper.schemas.PluginConfigsJsonSchemaConverter.PLUGIN_NAME_KEY;
+import static org.opensearch.dataprepper.schemas.PluginConfigsJsonSchemaConverter.PRIMARY_FIELDS_KEY;
 
 class PluginConfigsJsonSchemaConverterIT {
     static final String DEFAULT_PLUGINS_CLASSPATH = "org.opensearch.dataprepper.plugins";
@@ -45,8 +46,10 @@ void setUp() {
                         JakartaValidationOption.INCLUDE_PATTERN_EXPRESSIONS)
         );
         final PluginProvider pluginProvider = new ClasspathPluginProvider();
+        final PrimaryFieldsOverride primaryFieldsOverride = new PrimaryFieldsOverride();
         objectUnderTest = new PluginConfigsJsonSchemaConverter(
-                pluginProvider, new JsonSchemaConverter(modules, pluginProvider), TEST_URL, TEST_BASE_URL);
+                pluginProvider, new JsonSchemaConverter(modules, pluginProvider),
+                primaryFieldsOverride, TEST_URL, TEST_BASE_URL);
     }
 
     @ParameterizedTest
@@ -64,6 +67,7 @@ void testConvertPluginConfigsIntoJsonSchemas(final Class<?> pluginType) {
             }
             assertThat(schemaMap, notNullValue());
             assertThat(schemaMap.containsKey(PLUGIN_NAME_KEY), is(true));
+            assertThat(schemaMap.containsKey(PRIMARY_FIELDS_KEY), is(true));
             assertThat(((String) schemaMap.get(DOCUMENTATION_LINK_KEY)).startsWith(TEST_URL + TEST_BASE_URL),
                     is(true));
         });

data-prepper-plugin-schema/src/test/java/org/opensearch/dataprepper/schemas/PluginConfigsJsonSchemaConverterTest.java

@@ -14,6 +14,7 @@
 import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
 import org.opensearch.dataprepper.plugin.PluginProvider;
 
+import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
@@ -30,6 +31,7 @@
 import static org.opensearch.dataprepper.schemas.PluginConfigsJsonSchemaConverter.DOCUMENTATION_LINK_KEY;
 import static org.opensearch.dataprepper.schemas.PluginConfigsJsonSchemaConverter.PLUGIN_NAME_KEY;
 import static org.opensearch.dataprepper.schemas.PluginConfigsJsonSchemaConverter.PLUGIN_TYPE_NAME_TO_CLASS_MAP;
+import static org.opensearch.dataprepper.schemas.PluginConfigsJsonSchemaConverter.PRIMARY_FIELDS_KEY;
 
 @ExtendWith(MockitoExtension.class)
 class PluginConfigsJsonSchemaConverterTest {
@@ -39,6 +41,9 @@ class PluginConfigsJsonSchemaConverterTest {
     @Mock
     private JsonSchemaConverter jsonSchemaConverter;
 
+    @Mock
+    private PrimaryFieldsOverride primaryFieldsOverride;
+
     @Mock
     private PluginProvider pluginProvider;
 
@@ -79,6 +84,28 @@ void testConvertPluginConfigsIntoJsonSchemasHappyPath() throws JsonProcessingExc
         assertThat(schemaMap.get(DOCUMENTATION_LINK_KEY), equalTo(
                 "{{site.url}}{{site.baseurl}}/data-prepper/pipelines/configuration/null/test_plugin/"
         ));
+        assertThat(schemaMap.get(PRIMARY_FIELDS_KEY), equalTo(Collections.emptyList()));
+        assertThat(schemaMap.containsKey(PLUGIN_NAME_KEY), is(true));
+    }
+
+    @Test
+    void testConvertPluginConfigsIntoJsonSchemasWithPrimaryFieldsOverride() throws JsonProcessingException {
+        final String testPrimaryField = "test_field";
+        when(primaryFieldsOverride.getPrimaryFieldsForComponent(eq("test_plugin"))).thenReturn(
+                Set.of(testPrimaryField));
+        when(pluginProvider.findPluginClasses(eq(TestPluginType.class))).thenReturn(Set.of(TestPlugin.class));
+        final ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
+        when(jsonSchemaConverter.convertIntoJsonSchema(
+                any(SchemaVersion.class), any(OptionPreset.class), eq(TestPluginConfig.class))).thenReturn(objectNode);
+        final Map<String, String> result = objectUnderTest.convertPluginConfigsIntoJsonSchemas(
+                SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON, TestPluginType.class);
+        assertThat(result.size(), equalTo(1));
+        final Map<String, Object> schemaMap = OBJECT_MAPPER.readValue(result.get("test_plugin"), MAP_TYPE_REFERENCE);
+        assertThat(schemaMap, notNullValue());
+        assertThat(schemaMap.get(DOCUMENTATION_LINK_KEY), equalTo(
+                "{{site.url}}{{site.baseurl}}/data-prepper/pipelines/configuration/null/test_plugin/"
+        ));
+        assertThat(schemaMap.get(PRIMARY_FIELDS_KEY), equalTo(Collections.singletonList(testPrimaryField)));
         assertThat(schemaMap.containsKey(PLUGIN_NAME_KEY), is(true));
     }