diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginFactory.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginFactory.java index 183b726fe04400..bcf9108320842b 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginFactory.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/PluginFactory.java @@ -16,6 +16,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.function.BiFunction; import java.util.function.Function; @@ -48,17 +49,72 @@ public static PluginFactory merge( PluginFactory b, @Nullable BiFunction, PluginFactory> pluginFactoryProvider) { + + if (b.isEmpty()) { + return a; + } + if (a.isEmpty()) { + return b; + } + PluginConfiguration mergedPluginConfig = PluginConfiguration.merge(a.pluginConfiguration, b.pluginConfiguration); List mergedClassLoaders = Stream.concat(a.getClassLoaders().stream(), b.getClassLoaders().stream()) .collect(Collectors.toList()); - if (pluginFactoryProvider != null) { - return pluginFactoryProvider.apply(mergedPluginConfig, mergedClassLoaders); - } else { - return PluginFactory.withCustomClasspath(mergedPluginConfig, mergedClassLoaders); + if (!a.hasLoadedPlugins() && !b.hasLoadedPlugins()) { + if (pluginFactoryProvider != null) { + return pluginFactoryProvider.apply(mergedPluginConfig, mergedClassLoaders); + } else { + if (mergedPluginConfig + .streamAll() + .anyMatch(config -> config.getSpring() != null && config.getSpring().isEnabled())) { + throw new IllegalStateException( + "Unexpected Spring configuration found without a provided Spring Plugin Factory"); + } + return PluginFactory.withCustomClasspath(mergedPluginConfig, mergedClassLoaders); + } } + + PluginFactory loadedA = a.hasLoadedPlugins() ? a : a.loadPlugins(); + PluginFactory loadedB = b.hasLoadedPlugins() ? b : b.loadPlugins(); + + return new PluginFactory( + mergedPluginConfig, + mergedClassLoaders, + Stream.concat( + loadedA.aspectPayloadValidators.stream() + .filter( + aPlugin -> + loadedB.pluginConfiguration.getAspectPayloadValidators().stream() + .noneMatch(bConfig -> aPlugin.getConfig().isDisabledBy(bConfig))), + loadedB.aspectPayloadValidators.stream()) + .collect(Collectors.toList()), + Stream.concat( + loadedA.mutationHooks.stream() + .filter( + aPlugin -> + loadedB.pluginConfiguration.getMutationHooks().stream() + .noneMatch(bConfig -> aPlugin.getConfig().isDisabledBy(bConfig))), + loadedB.mutationHooks.stream()) + .collect(Collectors.toList()), + Stream.concat( + loadedA.mclSideEffects.stream() + .filter( + aPlugin -> + loadedB.pluginConfiguration.getMclSideEffects().stream() + .noneMatch(bConfig -> aPlugin.getConfig().isDisabledBy(bConfig))), + loadedB.mclSideEffects.stream()) + .collect(Collectors.toList()), + Stream.concat( + loadedA.mcpSideEffects.stream() + .filter( + aPlugin -> + loadedB.pluginConfiguration.getMcpSideEffects().stream() + .noneMatch(bConfig -> aPlugin.getConfig().isDisabledBy(bConfig))), + loadedB.mcpSideEffects.stream()) + .collect(Collectors.toList())); } @Getter private final PluginConfiguration pluginConfiguration; @@ -77,22 +133,62 @@ public PluginFactory( pluginConfiguration == null ? PluginConfiguration.EMPTY : pluginConfiguration; } + public PluginFactory( + @Nullable PluginConfiguration pluginConfiguration, + @Nonnull List classLoaders, + @Nonnull List aspectPayloadValidators, + @Nonnull List mutationHooks, + @Nonnull List mclSideEffects, + @Nonnull List mcpSideEffects) { + this.classLoaders = classLoaders; + this.pluginConfiguration = + pluginConfiguration == null ? PluginConfiguration.EMPTY : pluginConfiguration; + this.aspectPayloadValidators = applyDisable(aspectPayloadValidators); + this.mutationHooks = applyDisable(mutationHooks); + this.mclSideEffects = applyDisable(mclSideEffects); + this.mcpSideEffects = applyDisable(mcpSideEffects); + } + public PluginFactory loadPlugins() { - this.aspectPayloadValidators = buildAspectPayloadValidators(this.pluginConfiguration); - this.mutationHooks = buildMutationHooks(this.pluginConfiguration); - this.mclSideEffects = buildMCLSideEffects(this.pluginConfiguration); - this.mcpSideEffects = buildMCPSideEffects(this.pluginConfiguration); - logSummary( - Stream.of( - this.aspectPayloadValidators, - this.mutationHooks, - this.mclSideEffects, - this.mcpSideEffects) - .flatMap(List::stream) - .collect(Collectors.toList())); + if (this.aspectPayloadValidators != null + || this.mutationHooks != null + || this.mclSideEffects != null + || this.mcpSideEffects != null) { + log.error("Plugins are already loaded. Re-building plugins will be skipped."); + } else { + this.aspectPayloadValidators = buildAspectPayloadValidators(this.pluginConfiguration); + this.mutationHooks = buildMutationHooks(this.pluginConfiguration); + this.mclSideEffects = buildMCLSideEffects(this.pluginConfiguration); + this.mcpSideEffects = buildMCPSideEffects(this.pluginConfiguration); + logSummary( + Stream.of( + this.aspectPayloadValidators, + this.mutationHooks, + this.mclSideEffects, + this.mcpSideEffects) + .flatMap(List::stream) + .collect(Collectors.toList())); + } return this; } + public boolean isEmpty() { + return this.pluginConfiguration.isEmpty() + && Optional.ofNullable(this.aspectPayloadValidators).map(List::isEmpty).orElse(true) + && Optional.ofNullable(this.mutationHooks).map(List::isEmpty).orElse(true) + && Optional.ofNullable(this.mclSideEffects).map(List::isEmpty).orElse(true) + && Optional.ofNullable(this.mcpSideEffects).map(List::isEmpty).orElse(true); + } + + public boolean hasLoadedPlugins() { + return Stream.of( + this.aspectPayloadValidators, + this.mutationHooks, + this.mcpSideEffects, + this.mcpSideEffects) + .anyMatch(Objects::nonNull); + } + private void logSummary(List pluginSpecs) { if (!pluginSpecs.isEmpty()) { log.info( diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java index e9494c49a9efb2..5c086982557a33 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/plugins/config/PluginConfiguration.java @@ -27,6 +27,13 @@ public class PluginConfiguration { public static PluginConfiguration EMPTY = new PluginConfiguration(); + public boolean isEmpty() { + return aspectPayloadValidators.isEmpty() + && mutationHooks.isEmpty() + && mclSideEffects.isEmpty() + && mcpSideEffects.isEmpty(); + } + public static PluginConfiguration merge(PluginConfiguration a, PluginConfiguration b) { return new PluginConfiguration( Stream.concat( diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java index a9f903f4b7017d..cecf21849f3aaa 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/plugins/PluginsTest.java @@ -226,4 +226,60 @@ public void tripleMergeWithDisabled() throws EntityRegistryException { .count(), 0); } + + @Test + public void testEmptyMerges() throws EntityRegistryException { + ConfigEntityRegistry configEntityRegistry1 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_1)); + ConfigEntityRegistry emptyEntityRegistry = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_2), + (config, classLoaders) -> PluginFactory.empty()); + + MergedEntityRegistry mergedEntityRegistry = new MergedEntityRegistry(configEntityRegistry1); + mergedEntityRegistry.apply(emptyEntityRegistry); + assertEquals(mergedEntityRegistry.getPluginFactory(), configEntityRegistry1.getPluginFactory()); + + MergedEntityRegistry mergedEntityRegistry2 = new MergedEntityRegistry(emptyEntityRegistry); + mergedEntityRegistry2.apply(configEntityRegistry1); + assertEquals( + mergedEntityRegistry2.getPluginFactory(), configEntityRegistry1.getPluginFactory()); + } + + @Test + public void testUnloadedMerge() throws EntityRegistryException { + ConfigEntityRegistry configEntityRegistry1 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_1), + (config, classLoaders) -> new PluginFactory(config, classLoaders)); + ConfigEntityRegistry configEntityRegistry2 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_2), + (config, classLoaders) -> new PluginFactory(config, classLoaders)); + + MergedEntityRegistry mergedEntityRegistry = new MergedEntityRegistry(configEntityRegistry1); + mergedEntityRegistry.apply(configEntityRegistry2); + + assertEquals( + mergedEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + assertEquals( + mergedEntityRegistry.getAllMutationHooks().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + assertEquals( + mergedEntityRegistry.getAllMCLSideEffects().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + assertEquals( + mergedEntityRegistry.getAllMCPSideEffects().stream() + .filter(p -> p.getConfig().getSupportedOperations().contains("DELETE")) + .count(), + 1); + } } diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py index 347d0f88b0cd01..aac141ce310638 100644 --- a/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py +++ b/metadata-ingestion-modules/airflow-plugin/tests/integration/dags/snowflake_operator.py @@ -6,12 +6,21 @@ SNOWFLAKE_COST_TABLE = "costs" SNOWFLAKE_PROCESSED_TABLE = "processed_costs" + +def _fake_snowflake_execute(*args, **kwargs): + raise ValueError("mocked snowflake execute to not run queries") + + with DAG( "snowflake_operator", start_date=datetime(2023, 1, 1), schedule_interval=None, catchup=False, ) as dag: + # HACK: We don't want to send real requests to Snowflake. As a workaround, + # we can simply monkey-patch the operator. + SnowflakeOperator.execute = _fake_snowflake_execute # type: ignore + transform_cost_table = SnowflakeOperator( snowflake_conn_id="my_snowflake", task_id="transform_cost_table", diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_query.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_query.py index bbee4587e2b4df..9f655b34177fc6 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_query.py +++ b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_query.py @@ -685,12 +685,7 @@ def table_upstreams_with_column_lineage( t.query_start_time AS query_start_time, t.query_id AS query_id FROM - ( - SELECT * from snowflake.account_usage.access_history - WHERE - query_start_time >= to_timestamp_ltz({start_time_millis}, 3) - AND query_start_time < to_timestamp_ltz({end_time_millis}, 3) - ) t, + snowflake.account_usage.access_history t, lateral flatten(input => t.DIRECT_OBJECTS_ACCESSED) r, lateral flatten(input => t.OBJECTS_MODIFIED) w, lateral flatten(input => w.value : "columns", outer => true) wcols, @@ -780,12 +775,14 @@ def table_upstreams_with_column_lineage( queries AS ( select qid.downstream_table_name, qid.query_id, query_history.query_text, query_history.start_time from query_ids qid - JOIN ( + LEFT JOIN ( SELECT * FROM snowflake.account_usage.query_history WHERE query_history.start_time >= to_timestamp_ltz({start_time_millis}, 3) AND query_history.start_time < to_timestamp_ltz({end_time_millis}, 3) ) query_history on qid.query_id = query_history.query_id + WHERE qid.query_id is not null + AND query_history.query_text is not null ) SELECT h.downstream_table_name AS "DOWNSTREAM_TABLE_NAME", @@ -850,12 +847,7 @@ def table_upstreams_only( t.query_start_time AS query_start_time, t.query_id AS query_id FROM - ( - SELECT * from snowflake.account_usage.access_history - WHERE - query_start_time >= to_timestamp_ltz({start_time_millis}, 3) - AND query_start_time < to_timestamp_ltz({end_time_millis}, 3) - ) t, + snowflake.account_usage.access_history t, lateral flatten(input => t.DIRECT_OBJECTS_ACCESSED) r, lateral flatten(input => t.OBJECTS_MODIFIED) w WHERE diff --git a/metadata-service/plugin/build.gradle b/metadata-service/plugin/build.gradle index f519eba4921d2d..501589c24d60a2 100644 --- a/metadata-service/plugin/build.gradle +++ b/metadata-service/plugin/build.gradle @@ -13,14 +13,18 @@ dependencies { implementation externalDependency.jacksonDataFormatYaml implementation externalDependency.jacksonJDK8 implementation externalDependency.jacksonDataPropertyFormat - implementation externalDependency.logbackClassic; + implementation externalDependency.logbackClassic implementation externalDependency.slf4jApi compileOnly externalDependency.lombok + annotationProcessor externalDependency.lombok + testImplementation project(':test-models') + testImplementation project(path: ':test-models', configuration: 'testDataTemplate') testImplementation externalDependency.mockito testImplementation externalDependency.testng - annotationProcessor externalDependency.lombok + testCompileOnly externalDependency.lombok + testAnnotationProcessor externalDependency.lombok } test { diff --git a/metadata-service/plugin/src/main/java/com/datahub/plugins/metadata/aspect/SpringPluginFactory.java b/metadata-service/plugin/src/main/java/com/datahub/plugins/metadata/aspect/SpringPluginFactory.java index 3953ab8a456363..dcedbec50b7938 100644 --- a/metadata-service/plugin/src/main/java/com/datahub/plugins/metadata/aspect/SpringPluginFactory.java +++ b/metadata-service/plugin/src/main/java/com/datahub/plugins/metadata/aspect/SpringPluginFactory.java @@ -106,16 +106,27 @@ protected List build( try { Class clazz = classLoader.loadClass(config.getClassName()); - final T plugin; + final List plugins; if (config.getSpring().getName() == null) { - plugin = (T) springApplicationContext.getBean(clazz); + plugins = + springApplicationContext.getBeansOfType(clazz).values().stream() + .map(plugin -> (T) plugin) + .collect(Collectors.toList()); } else { - plugin = (T) springApplicationContext.getBean(config.getSpring().getName(), clazz); + plugins = + List.of((T) springApplicationContext.getBean(config.getSpring().getName(), clazz)); } - if (plugin.enabled()) { - result.add((T) plugin.setConfig(config)); - } + plugins.stream() + .filter(plugin -> plugin.enabled()) + .forEach( + plugin -> { + if (plugin.getConfig() != null) { + result.add(plugin); + } else { + result.add((T) plugin.setConfig(config)); + } + }); loaded = true; break; diff --git a/metadata-service/plugin/src/test/java/com/datahub/plugins/metadata/aspect/SpringPluginFactoryTest.java b/metadata-service/plugin/src/test/java/com/datahub/plugins/metadata/aspect/SpringPluginFactoryTest.java new file mode 100644 index 00000000000000..d321697304afc4 --- /dev/null +++ b/metadata-service/plugin/src/test/java/com/datahub/plugins/metadata/aspect/SpringPluginFactoryTest.java @@ -0,0 +1,188 @@ +package com.datahub.plugins.metadata.aspect; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + +import com.datahub.test.TestEntityProfile; +import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.RetrieverContext; +import com.linkedin.metadata.aspect.batch.BatchItem; +import com.linkedin.metadata.aspect.batch.ChangeMCP; +import com.linkedin.metadata.aspect.batch.MCLItem; +import com.linkedin.metadata.aspect.batch.MCPItem; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import com.linkedin.metadata.aspect.plugins.hooks.MCLSideEffect; +import com.linkedin.metadata.aspect.plugins.hooks.MCPSideEffect; +import com.linkedin.metadata.aspect.plugins.hooks.MutationHook; +import com.linkedin.metadata.aspect.plugins.validation.AspectPayloadValidator; +import com.linkedin.metadata.aspect.plugins.validation.AspectValidationException; +import com.linkedin.metadata.models.EntitySpec; +import com.linkedin.metadata.models.EventSpec; +import com.linkedin.metadata.models.registry.ConfigEntityRegistry; +import com.linkedin.metadata.models.registry.EntityRegistryException; +import com.linkedin.metadata.models.registry.MergedEntityRegistry; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.stream.Stream; +import javax.annotation.Nonnull; +import lombok.Getter; +import lombok.Setter; +import lombok.experimental.Accessors; +import org.springframework.context.annotation.Configuration; +import org.testng.annotations.BeforeTest; +import org.testng.annotations.Test; + +@Configuration +public class SpringPluginFactoryTest { + + public static String REGISTRY_FILE_1 = "test-entity-registry-plugins-1.yml"; + public static String REGISTRY_SPRING_FILE_1 = "test-entity-registry-spring-plugins-1.yml"; + + @BeforeTest + public void disableAssert() { + PathSpecBasedSchemaAnnotationVisitor.class + .getClassLoader() + .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false); + } + + @Test + public void testMergedEntityRegistryWithSpringPluginFactory() throws EntityRegistryException { + ConfigEntityRegistry configEntityRegistry1 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_FILE_1)); + ConfigEntityRegistry configEntityRegistry2 = + new ConfigEntityRegistry( + TestEntityProfile.class.getClassLoader().getResourceAsStream(REGISTRY_SPRING_FILE_1), + (config, classLoaders) -> + new SpringPluginFactory( + null, config, List.of(SpringPluginFactoryTest.class.getClassLoader()))); + + MergedEntityRegistry mergedEntityRegistry = new MergedEntityRegistry(configEntityRegistry1); + mergedEntityRegistry.apply(configEntityRegistry2); + + Map entitySpecs = mergedEntityRegistry.getEntitySpecs(); + Map eventSpecs = mergedEntityRegistry.getEventSpecs(); + assertEquals(entitySpecs.values().size(), 2); + assertEquals(eventSpecs.values().size(), 1); + + EntitySpec entitySpec = mergedEntityRegistry.getEntitySpec("dataset"); + assertEquals(entitySpec.getName(), "dataset"); + assertEquals(entitySpec.getKeyAspectSpec().getName(), "datasetKey"); + assertEquals(entitySpec.getAspectSpecs().size(), 4); + assertNotNull(entitySpec.getAspectSpec("datasetKey")); + assertNotNull(entitySpec.getAspectSpec("datasetProperties")); + assertNotNull(entitySpec.getAspectSpec("schemaMetadata")); + assertNotNull(entitySpec.getAspectSpec("status")); + + entitySpec = mergedEntityRegistry.getEntitySpec("chart"); + assertEquals(entitySpec.getName(), "chart"); + assertEquals(entitySpec.getKeyAspectSpec().getName(), "chartKey"); + assertEquals(entitySpec.getAspectSpecs().size(), 3); + assertNotNull(entitySpec.getAspectSpec("chartKey")); + assertNotNull(entitySpec.getAspectSpec("chartInfo")); + assertNotNull(entitySpec.getAspectSpec("status")); + + EventSpec eventSpec = mergedEntityRegistry.getEventSpec("testEvent"); + assertEquals(eventSpec.getName(), "testEvent"); + assertNotNull(eventSpec.getPegasusSchema()); + + assertEquals( + mergedEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "chart", "status")) + .count(), + 2); + assertEquals( + mergedEntityRegistry.getAllAspectPayloadValidators().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "chart", "status")) + .count(), + 1); + + assertEquals( + mergedEntityRegistry.getAllMCPSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "dataset", "datasetKey")) + .count(), + 2); + assertEquals( + mergedEntityRegistry.getAllMCPSideEffects().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "dataset", "datasetKey")) + .count(), + 1); + + assertEquals( + mergedEntityRegistry.getAllMutationHooks().stream() + .filter(validator -> validator.shouldApply(ChangeType.UPSERT, "*", "schemaMetadata")) + .count(), + 2); + assertEquals( + mergedEntityRegistry.getAllMutationHooks().stream() + .filter(validator -> validator.shouldApply(ChangeType.DELETE, "*", "schemaMetadata")) + .count(), + 1); + } + + /* + * Various test plugins to be injected with Spring + */ + @Getter + @Setter + @Accessors(chain = true) + public static class TestValidator extends AspectPayloadValidator { + + public AspectPluginConfig config; + + @Override + protected Stream validateProposedAspects( + @Nonnull Collection mcpItems, + @Nonnull RetrieverContext retrieverContext) { + return mcpItems.stream().map(i -> AspectValidationException.forItem(i, "test error")); + } + + @Override + protected Stream validatePreCommitAspects( + @Nonnull Collection changeMCPs, @Nonnull RetrieverContext retrieverContext) { + return Stream.empty(); + } + } + + @Getter + @Setter + @Accessors(chain = true) + public static class TestMutator extends MutationHook { + public AspectPluginConfig config; + } + + @Getter + @Setter + @Accessors(chain = true) + public static class TestMCPSideEffect extends MCPSideEffect { + + public AspectPluginConfig config; + + @Override + protected Stream applyMCPSideEffect( + Collection changeMCPS, @Nonnull RetrieverContext retrieverContext) { + return changeMCPS.stream(); + } + + @Override + protected Stream postMCPSideEffect( + Collection mclItems, @Nonnull RetrieverContext retrieverContext) { + return Stream.of(); + } + } + + @Getter + @Setter + @Accessors(chain = true) + public static class TestMCLSideEffect extends MCLSideEffect { + public AspectPluginConfig config; + + @Override + protected Stream applyMCLSideEffect( + @Nonnull Collection batchItems, @Nonnull RetrieverContext retrieverContext) { + return null; + } + } +} diff --git a/metadata-service/plugin/src/test/java/com/datahub/plugins/metadata/aspect/test/TestSpringPluginConfiguration.java b/metadata-service/plugin/src/test/java/com/datahub/plugins/metadata/aspect/test/TestSpringPluginConfiguration.java new file mode 100644 index 00000000000000..589ae614e2d2f3 --- /dev/null +++ b/metadata-service/plugin/src/test/java/com/datahub/plugins/metadata/aspect/test/TestSpringPluginConfiguration.java @@ -0,0 +1,109 @@ +package com.datahub.plugins.metadata.aspect.test; + +import com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest; +import com.linkedin.events.metadata.ChangeType; +import com.linkedin.metadata.aspect.plugins.config.AspectPluginConfig; +import java.util.List; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +/** This is the Spring class used for the test */ +@Configuration +public class TestSpringPluginConfiguration { + @Bean + public SpringPluginFactoryTest.TestValidator springValidator1() { + SpringPluginFactoryTest.TestValidator testValidator = + new SpringPluginFactoryTest.TestValidator(); + testValidator.setConfig( + AspectPluginConfig.builder() + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("dataset") + .aspectName("status") + .build())) + .className(SpringPluginFactoryTest.TestValidator.class.getName()) + .enabled(true) + .supportedOperations(List.of(ChangeType.UPSERT.toString())) + .build()); + return testValidator; + } + + @Bean + public SpringPluginFactoryTest.TestValidator springValidator2() { + SpringPluginFactoryTest.TestValidator testValidator = + new SpringPluginFactoryTest.TestValidator(); + testValidator.setConfig( + AspectPluginConfig.builder() + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("*") + .aspectName("status") + .build())) + .className(SpringPluginFactoryTest.TestValidator.class.getName()) + .enabled(true) + .supportedOperations(List.of(ChangeType.DELETE.toString())) + .build()); + return testValidator; + } + + @Bean + public SpringPluginFactoryTest.TestMutator springMutator() { + SpringPluginFactoryTest.TestMutator testMutator = new SpringPluginFactoryTest.TestMutator(); + testMutator.setConfig( + AspectPluginConfig.builder() + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("*") + .aspectName("schemaMetadata") + .build())) + .className(SpringPluginFactoryTest.TestMutator.class.getName()) + .enabled(true) + .supportedOperations( + List.of(ChangeType.UPSERT.toString(), ChangeType.DELETE.toString())) + .build()); + return testMutator; + } + + @Bean + public SpringPluginFactoryTest.TestMCPSideEffect springMCPSideEffect() { + SpringPluginFactoryTest.TestMCPSideEffect testMCPSideEffect = + new SpringPluginFactoryTest.TestMCPSideEffect(); + testMCPSideEffect.setConfig( + AspectPluginConfig.builder() + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("dataset") + .aspectName("datasetKey") + .build())) + .className(SpringPluginFactoryTest.TestMCPSideEffect.class.getName()) + .enabled(true) + .supportedOperations( + List.of(ChangeType.UPSERT.toString(), ChangeType.DELETE.toString())) + .build()); + return testMCPSideEffect; + } + + @Bean + public SpringPluginFactoryTest.TestMCLSideEffect springMCLSideEffect() { + SpringPluginFactoryTest.TestMCLSideEffect testMCLSideEffect = + new SpringPluginFactoryTest.TestMCLSideEffect(); + testMCLSideEffect.setConfig( + AspectPluginConfig.builder() + .supportedEntityAspectNames( + List.of( + AspectPluginConfig.EntityAspectName.builder() + .entityName("chart") + .aspectName("ChartInfo") + .build())) + .className(SpringPluginFactoryTest.TestMCLSideEffect.class.getName()) + .enabled(true) + .supportedOperations( + List.of(ChangeType.UPSERT.toString(), ChangeType.DELETE.toString())) + .build()); + return testMCLSideEffect; + } +} diff --git a/metadata-service/plugin/src/test/resources/test-entity-registry-plugins-1.yml b/metadata-service/plugin/src/test/resources/test-entity-registry-plugins-1.yml new file mode 100644 index 00000000000000..903e32b623db0f --- /dev/null +++ b/metadata-service/plugin/src/test/resources/test-entity-registry-plugins-1.yml @@ -0,0 +1,67 @@ +id: test-registry-1 +entities: + - name: dataset + keyAspect: datasetKey + category: core + aspects: + - datasetProperties + - schemaMetadata + - status + - name: chart + keyAspect: chartKey + aspects: + - chartInfo + - status +events: + - name: testEvent + +plugins: + aspectPayloadValidators: + # All status aspects on any entity + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestValidator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: '*' + aspectName: status + # Chart status only + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestValidator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: chart + aspectName: status + # Disabled + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestValidator' + enabled: false + supportedOperations: + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: status + mutationHooks: + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestMutator' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: '*' + aspectName: schemaMetadata + mclSideEffects: + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestMCLSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: chart + aspectName: chartInfo + mcpSideEffects: + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestMCPSideEffect' + enabled: true + supportedOperations: + - UPSERT + supportedEntityAspectNames: + - entityName: dataset + aspectName: datasetKey diff --git a/metadata-service/plugin/src/test/resources/test-entity-registry-spring-plugins-1.yml b/metadata-service/plugin/src/test/resources/test-entity-registry-spring-plugins-1.yml new file mode 100644 index 00000000000000..a7ae97e9bceb65 --- /dev/null +++ b/metadata-service/plugin/src/test/resources/test-entity-registry-spring-plugins-1.yml @@ -0,0 +1,52 @@ +id: test-registry-spring-1 +entities: [] +plugins: + aspectPayloadValidators: + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestValidator' + packageScan: + - com.datahub.plugins.metadata.aspect.test + spring: + enabled: true + enabled: true + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: status + mutationHooks: + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestMutator' + packageScan: + - com.datahub.plugins.metadata.aspect.test + spring: + enabled: true + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: '*' + aspectName: schemaMetadata + mclSideEffects: + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestMCLSideEffect' + packageScan: + - com.datahub.plugins.metadata.aspect.test + spring: + enabled: true + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: chart + aspectName: chartInfo + mcpSideEffects: + - className: 'com.datahub.plugins.metadata.aspect.SpringPluginFactoryTest$TestMCPSideEffect' + packageScan: + - com.datahub.plugins.metadata.aspect.test + spring: + enabled: true + supportedOperations: + - UPSERT + - DELETE + supportedEntityAspectNames: + - entityName: dataset + aspectName: datasetKey