diff --git a/bin/distributed-test/scripts/trigger-reindex.sh b/bin/distributed-test/scripts/trigger-reindex.sh index fa39548ed27c..609b70665b65 100755 --- a/bin/distributed-test/scripts/trigger-reindex.sh +++ b/bin/distributed-test/scripts/trigger-reindex.sh @@ -8,6 +8,7 @@ PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" # Default values SERVER_URL="http://localhost:8585" +RECREATE_INDEX=false ENTITY_TYPES="" BATCH_SIZE=100 PARTITION_SIZE=10000 @@ -19,6 +20,10 @@ while [[ $# -gt 0 ]]; do SERVER_URL="$2" shift 2 ;; + --recreate) + RECREATE_INDEX=true + shift + ;; --entities) ENTITY_TYPES="$2" shift 2 @@ -36,6 +41,7 @@ while [[ $# -gt 0 ]]; do echo "" echo "Options:" echo " --server URL Target server URL (default: http://localhost:8585)" + echo " --recreate Drop and recreate indices before reindexing" echo " --entities TYPES Comma-separated entity types to reindex (default: all)" echo " --batch-size NUM Batch size for indexing (default: 100)" echo " --partition-size NUM Partition size for distributed indexing (default: 10000, range: 1000-50000)" @@ -45,6 +51,7 @@ while [[ $# -gt 0 ]]; do echo "Examples:" echo " $0 # Reindex all on server 1" echo " $0 --server http://localhost:8587 # Trigger on server 2" + echo " $0 --recreate # Drop and recreate indices" echo " $0 --entities table,dashboard # Reindex only tables and dashboards" echo " $0 --partition-size 2000 # Use smaller partitions for better distribution" exit 0 @@ -60,7 +67,7 @@ echo "======================================" echo "Triggering Search Reindexing" echo "======================================" echo "Server: $SERVER_URL" -echo "Indexing mode: staged indexes with alias promotion" +echo "Recreate indices: $RECREATE_INDEX" echo "Batch size: $BATCH_SIZE" echo "Partition size: $PARTITION_SIZE" if [ -n "$ENTITY_TYPES" ]; then @@ -89,6 +96,13 @@ fi echo "Authenticated successfully." echo "" +# Build the reindex request body +if [ "$RECREATE_INDEX" == "true" ]; then + RECREATE_FLAG="true" +else + RECREATE_FLAG="false" +fi + # Build entities array if [ -n "$ENTITY_TYPES" ]; then # Convert comma-separated to JSON array @@ -99,9 +113,11 @@ fi REQUEST_BODY=$(cat <The read-only tests ({@link #analyzeReturnsRecommendationsForKnownTables}, {@link - * #dryRunDoesNotMutateReloptions}) run against the real catalog tables that the IT bootstrap - * created via migrations. - * - *

Tests that exercise the write path ({@link #applyExecutesAndIsIdempotent}, {@link - * #analyzeOneRunsOnIsolatedTable}) deliberately use a private throwaway table — never a real - * catalog table. Reason: {@code ALTER TABLE} on a shared production table bumps MySQL's per-table - * metadata version, which invalidates JDBC prepared-statement caches across the whole - * Testcontainer. When that table has a {@code JSON} column (e.g. {@code entity_relationship}), the - * driver's re-prepared metadata sometimes returns the column type as {@code VARBINARY}, and - * subsequent {@code INSERT} statements fail with {@code "Cannot create a JSON value from a string - * with CHARACTER SET 'binary'"}. We saw this break {@code GlossaryTermRelationsIT}, - * {@code DomainResourceIT}, and the lineage ITs in CI when an earlier version of this test applied - * settings to {@code entity_relationship}. The recommendations themselves are sound — the IT just - * cannot afford the side effect on a shared DB. - * - *

Sequential because {@code @BeforeEach} / {@code @AfterEach} create and drop the same isolated - * table by name; concurrent execution would race. - */ -@Execution(ExecutionMode.SAME_THREAD) -class DbTuneIT { - - /** Table created and dropped per test — never a catalog table. Safe blast radius. */ - private static final String ISOLATED_TABLE = "dbtune_it_isolated_table"; - - /** A real catalog table used only by the read-only tests to assert against the live schema. */ - private static final String READ_ONLY_PROBE_TABLE = "entity_relationship"; - - @BeforeEach - void createIsolatedTable() { - Jdbi jdbi = TestSuiteBootstrap.getJdbi(); - ConnectionType connType = currentConnectionType(); - jdbi.useHandle( - handle -> { - handle.execute("DROP TABLE IF EXISTS " + quoteIdent(connType, ISOLATED_TABLE)); - if (connType == ConnectionType.POSTGRES) { - handle.execute( - "CREATE TABLE " + quoteIdent(connType, ISOLATED_TABLE) + " (id INT PRIMARY KEY)"); - } else { - handle.execute( - "CREATE TABLE " - + quoteIdent(connType, ISOLATED_TABLE) - + " (id INT PRIMARY KEY) ENGINE=InnoDB"); - } - }); - } - - @AfterEach - void dropIsolatedTable() { - Jdbi jdbi = TestSuiteBootstrap.getJdbi(); - ConnectionType connType = currentConnectionType(); - jdbi.useHandle( - handle -> handle.execute("DROP TABLE IF EXISTS " + quoteIdent(connType, ISOLATED_TABLE))); - } - - @Test - void analyzeReturnsRecommendationsForKnownTables() { - AutoTuner tuner = currentTuner(); - Jdbi jdbi = TestSuiteBootstrap.getJdbi(); - - DbTuneResult result = jdbi.withHandle(tuner::analyze); - - assertNotNull(result); - assertNotNull(result.engineVersion()); - assertFalse(result.tableRecommendations().isEmpty(), "Expected at least one recommendation"); - assertTrue( - result.tableRecommendations().stream() - .anyMatch(r -> READ_ONLY_PROBE_TABLE.equals(r.tableName())), - READ_ONLY_PROBE_TABLE + " should be in the recommendations"); - } - - @Test - void applyExecutesAndIsIdempotent() { - AutoTuner tuner = currentTuner(); - Jdbi jdbi = TestSuiteBootstrap.getJdbi(); - ConnectionType connType = currentConnectionType(); - TableRecommendation rec = recommendationForIsolatedTable(connType); - - String built = tuner.buildAlterStatement(rec); - assertTrue(built.contains(ISOLATED_TABLE), "ALTER target table mismatch: " + built); - - jdbi.useHandle(handle -> tuner.apply(handle, rec)); - Map after = - jdbi.withHandle(handle -> tuner.currentSettingsForTable(handle, ISOLATED_TABLE)); - assertSettingsPersisted(rec.recommendedSettings(), after); - - // Apply a second time — must be idempotent (no exception, no value drift). - jdbi.useHandle(handle -> tuner.apply(handle, rec)); - Map afterSecond = - jdbi.withHandle(handle -> tuner.currentSettingsForTable(handle, ISOLATED_TABLE)); - assertEquals(after, afterSecond, "Apply should be idempotent"); - } - - private void assertSettingsPersisted( - final Map expected, final Map actual) { - for (Map.Entry e : expected.entrySet()) { - String key = e.getKey(); - // Postgres lowercases reloption keys; MySQL uppercases STATS_*. Look up case-insensitively. - String got = - actual.entrySet().stream() - .filter(a -> a.getKey().equalsIgnoreCase(key)) - .map(Map.Entry::getValue) - .findFirst() - .orElse(null); - assertNotNull(got, "Missing setting after apply: " + key + " (got " + actual + ")"); - assertEquals( - Double.parseDouble(e.getValue()), - Double.parseDouble(got), - 0.0, - "Setting " + key + " did not take effect: expected " + e.getValue() + ", got " + got); - } - } - - @Test - void analyzeOneRunsOnIsolatedTable() { - AutoTuner tuner = currentTuner(); - Jdbi jdbi = TestSuiteBootstrap.getJdbi(); - - jdbi.useHandle(handle -> tuner.analyzeOne(handle, ISOLATED_TABLE)); - } - - @Test - void diagnoseCompletesWithoutErrorAndReturnsStructuredResult() { - Diagnostic diagnostic = currentDiagnostic(); - Jdbi jdbi = TestSuiteBootstrap.getJdbi(); - - DbTuneDiagnosis diagnosis = jdbi.withHandle(diagnostic::diagnose); - - assertNotNull(diagnosis, "diagnose() must return a non-null diagnosis"); - assertNotNull(diagnosis.findings(), "findings list must be present (empty allowed)"); - assertNotNull(diagnosis.notes(), "notes list must be present (empty allowed)"); - // On a freshly-bootstrapped IT DB we expect either: - // - an empty diagnosis (nothing has accumulated yet to flag), OR - // - notes about missing optional extensions like pg_stat_statements. - // Either is fine — what we're really asserting is the diagnostic ran end-to-end without - // throwing on the live schema. - } - - @Test - void dryRunDoesNotMutateReloptions() { - AutoTuner tuner = currentTuner(); - Jdbi jdbi = TestSuiteBootstrap.getJdbi(); - - Map before = currentSettingsFor(tuner, jdbi, READ_ONLY_PROBE_TABLE); - - DbTuneResult result = jdbi.withHandle(tuner::analyze); - assertNotNull(result); - - Map after = currentSettingsFor(tuner, jdbi, READ_ONLY_PROBE_TABLE); - assertEquals(before, after, "Analyze (dry-run) must not change table settings"); - } - - // ---- helpers ---- - - private AutoTuner currentTuner() { - return currentConnectionType() == ConnectionType.POSTGRES - ? new PostgresAutoTuner() - : new MysqlAutoTuner(); - } - - private Diagnostic currentDiagnostic() { - return currentConnectionType() == ConnectionType.POSTGRES - ? new PostgresDiagnostic() - : new MysqlDiagnostic(); - } - - private ConnectionType currentConnectionType() { - return "mysql".equalsIgnoreCase(System.getProperty("databaseType", "postgres")) - ? ConnectionType.MYSQL - : ConnectionType.POSTGRES; - } - - /** - * Builds a {@link TableRecommendation} pointing at {@link #ISOLATED_TABLE} with engine-appropriate - * settings. We construct it directly rather than going through {@code analyze()} because the - * isolated table is intentionally NOT in the static catalog — that's how we keep the apply path - * off shared production tables. - */ - private TableRecommendation recommendationForIsolatedTable(final ConnectionType connType) { - Map recommended = - connType == ConnectionType.POSTGRES - ? Map.of("autovacuum_vacuum_scale_factor", "0.05") - : Map.of("STATS_PERSISTENT", "1", "STATS_AUTO_RECALC", "1"); - return new TableRecommendation( - ISOLATED_TABLE, Action.APPLY, 0L, 0L, Map.of(), recommended, "Isolated IT test table"); - } - - /** - * Re-runs analyze and projects out the {@link TableRecommendation#currentSettings()} for the - * named table. Going through the same code path that built the original recommendation keeps the - * assertion stable across either dialect's parsing rules. - */ - private Map currentSettingsFor( - final AutoTuner tuner, final Jdbi jdbi, final String tableName) { - return jdbi.withHandle(tuner::analyze).tableRecommendations().stream() - .filter(r -> tableName.equals(r.tableName())) - .findFirst() - .map(TableRecommendation::currentSettings) - .orElse(Map.of()); - } - - private static String quoteIdent(final ConnectionType connType, final String identifier) { - if (!identifier.matches("[a-zA-Z_][a-zA-Z0-9_]*")) { - throw new IllegalArgumentException("Refusing unsafe identifier: " + identifier); - } - return connType == ConnectionType.POSTGRES ? "\"" + identifier + "\"" : "`" + identifier + "`"; - } -} diff --git a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/RdfResourceIT.java b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/RdfResourceIT.java index 99d618e35c53..51b126d19cb6 100644 --- a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/RdfResourceIT.java +++ b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/RdfResourceIT.java @@ -1,8 +1,15 @@ package org.openmetadata.it.tests; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assumptions.assumeTrue; +import java.net.URI; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.List; import org.awaitility.Awaitility; @@ -15,6 +22,7 @@ import org.openmetadata.it.factories.DatabaseSchemaTestFactory; import org.openmetadata.it.factories.DatabaseServiceTestFactory; import org.openmetadata.it.util.RdfTestUtils; +import org.openmetadata.it.util.SdkClients; import org.openmetadata.it.util.TestNamespace; import org.openmetadata.it.util.TestNamespaceExtension; import org.openmetadata.schema.api.configuration.rdf.RdfConfiguration; @@ -23,6 +31,7 @@ import org.openmetadata.schema.entity.data.Table; import org.openmetadata.schema.entity.services.DatabaseService; import org.openmetadata.schema.type.Column; +import org.openmetadata.schema.type.TableConstraint; import org.openmetadata.sdk.fluent.Tables; import org.openmetadata.sdk.fluent.builders.ColumnBuilder; import org.openmetadata.service.rdf.RdfUpdater; @@ -43,6 +52,10 @@ public class RdfResourceIT { private static final String TABLE_RDF_TYPE = "dcat:Dataset"; + private static final String BASE_URI = "https://open-metadata.org/"; + private static final String OM_NS = BASE_URI + "ontology/"; + private static final HttpClient HTTP_CLIENT = + HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(30)).build(); @BeforeAll static void enableRdf() { @@ -51,9 +64,9 @@ static void enableRdf() { "RDF is disabled for this test run. Use the RDF test profile to execute RdfResourceIT."); RdfConfiguration rdfConfig = new RdfConfiguration(); rdfConfig.setEnabled(true); - rdfConfig.setBaseUri(java.net.URI.create("https://open-metadata.org/")); + rdfConfig.setBaseUri(URI.create(BASE_URI)); rdfConfig.setStorageType(RdfConfiguration.StorageType.FUSEKI); - rdfConfig.setRemoteEndpoint(java.net.URI.create(TestSuiteBootstrap.getFusekiEndpoint())); + rdfConfig.setRemoteEndpoint(URI.create(TestSuiteBootstrap.getFusekiEndpoint())); rdfConfig.setUsername("admin"); rdfConfig.setPassword("test-admin"); rdfConfig.setDataset("openmetadata"); @@ -179,4 +192,251 @@ void testEntityUpdateInRdf(TestNamespace ns) { .pollInterval(Duration.ofMillis(500)) .untilAsserted(() -> RdfTestUtils.verifyEntityUpdatedInRdf(updated)); } + + // --------------------------------------------------------------------------- + // Phase-1 knowledge-graph fidelity tests (P1.1, P1.2, P1.7, P1.9). + // These exercise the new column / constraint / SHACL / ontology behavior + // against the same Fuseki container used by the tests above. + // --------------------------------------------------------------------------- + + @Test + void testColumnIsNamedRdfResource(TestNamespace ns) { + DatabaseService service = DatabaseServiceTestFactory.createPostgres(ns); + DatabaseSchema schema = DatabaseSchemaTestFactory.createSimple(ns, service); + + CreateTable createRequest = new CreateTable(); + createRequest.setName(ns.prefix("rdfColumnUriTable")); + createRequest.setDatabaseSchema(schema.getFullyQualifiedName()); + createRequest.setColumns( + List.of( + ColumnBuilder.of("id", "BIGINT").primaryKey().build(), + ColumnBuilder.of("email", "VARCHAR").dataLength(255).unique().build())); + + Table table = Tables.create(createRequest); + assertNotNull(table.getId()); + + String idColumnFqn = table.getFullyQualifiedName() + ".id"; + String emailColumnFqn = table.getFullyQualifiedName() + ".email"; + + Awaitility.await() + .atMost(Duration.ofSeconds(15)) + .pollInterval(Duration.ofMillis(500)) + .untilAsserted( + () -> { + assertTrue( + columnIsTypedAsOmColumn(idColumnFqn), + "PK column should be a named om:Column resource at the FQN-derived URI"); + assertTrue( + columnIsTypedAsOmColumn(emailColumnFqn), + "UNIQUE column should be a named om:Column resource at the FQN-derived URI"); + }); + } + + @Test + void testColumnConstraintFlagsInRdf(TestNamespace ns) { + DatabaseService service = DatabaseServiceTestFactory.createPostgres(ns); + DatabaseSchema schema = DatabaseSchemaTestFactory.createSimple(ns, service); + + CreateTable createRequest = new CreateTable(); + createRequest.setName(ns.prefix("rdfConstraintTable")); + createRequest.setDatabaseSchema(schema.getFullyQualifiedName()); + createRequest.setColumns( + List.of( + ColumnBuilder.of("id", "BIGINT").primaryKey().build(), + ColumnBuilder.of("email", "VARCHAR").dataLength(255).unique().build(), + ColumnBuilder.of("country", "VARCHAR").dataLength(2).notNull().build())); + + Table table = Tables.create(createRequest); + String idFqn = table.getFullyQualifiedName() + ".id"; + String emailFqn = table.getFullyQualifiedName() + ".email"; + String countryFqn = table.getFullyQualifiedName() + ".country"; + + Awaitility.await() + .atMost(Duration.ofSeconds(15)) + .pollInterval(Duration.ofMillis(500)) + .untilAsserted( + () -> { + assertTrue( + columnHasBooleanProperty(idFqn, "isPrimaryKey", true), + "Primary-key column should set om:isPrimaryKey true"); + assertTrue( + columnHasBooleanProperty(idFqn, "isNullable", false), + "Primary-key column should set om:isNullable false"); + assertTrue( + columnHasBooleanProperty(idFqn, "isUnique", true), + "Primary-key column should also set om:isUnique true (PKs are unique)"); + assertTrue( + columnHasBooleanProperty(emailFqn, "isUnique", true), + "UNIQUE column should set om:isUnique true"); + assertTrue( + columnHasBooleanProperty(countryFqn, "isNullable", false), + "NOT_NULL column should set om:isNullable false"); + }); + } + + @Test + void testForeignKeyReferencesInRdf(TestNamespace ns) { + DatabaseService service = DatabaseServiceTestFactory.createPostgres(ns); + DatabaseSchema schema = DatabaseSchemaTestFactory.createSimple(ns, service); + + CreateTable customers = new CreateTable(); + customers.setName(ns.prefix("rdfFkCustomers")); + customers.setDatabaseSchema(schema.getFullyQualifiedName()); + customers.setColumns( + List.of( + ColumnBuilder.of("id", "BIGINT").primaryKey().build(), + ColumnBuilder.of("name", "VARCHAR").dataLength(255).build())); + Table customersTable = Tables.create(customers); + + String customerIdFqn = customersTable.getFullyQualifiedName() + ".id"; + + CreateTable orders = new CreateTable(); + orders.setName(ns.prefix("rdfFkOrders")); + orders.setDatabaseSchema(schema.getFullyQualifiedName()); + orders.setColumns( + List.of( + ColumnBuilder.of("order_id", "BIGINT").primaryKey().build(), + ColumnBuilder.of("customer_id", "BIGINT").notNull().build())); + + TableConstraint fk = + new TableConstraint() + .withConstraintType(TableConstraint.ConstraintType.FOREIGN_KEY) + .withColumns(List.of("customer_id")) + .withReferredColumns(List.of(customerIdFqn)) + .withRelationshipType(TableConstraint.RelationshipType.MANY_TO_ONE); + orders.setTableConstraints(List.of(fk)); + + Table ordersTable = Tables.create(orders); + + String orderCustomerIdFqn = ordersTable.getFullyQualifiedName() + ".customer_id"; + + Awaitility.await() + .atMost(Duration.ofSeconds(15)) + .pollInterval(Duration.ofMillis(500)) + .untilAsserted( + () -> { + assertTrue( + columnReferencesColumn(orderCustomerIdFqn, customerIdFqn), + "FOREIGN_KEY constraint should produce direct om:references triple between source and referred column"); + assertTrue( + tableHasConstraintOfType(ordersTable.getFullyQualifiedName(), "FOREIGN_KEY"), + "Table should declare a FOREIGN_KEY om:TableConstraint resource"); + }); + } + + @Test + void testOntologyEndpointServesBumpedVersion() throws Exception { + String url = SdkClients.getServerUrl() + "/v1/rdf/ontology"; + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(url)) + .header("Authorization", "Bearer " + SdkClients.getAdminToken()) + .header("Accept", "text/turtle") + .GET() + .build(); + HttpResponse response = HTTP_CLIENT.send(request, HttpResponse.BodyHandlers.ofString()); + assertNotNull(response); + assertTrue( + response.statusCode() == 200, + "GET /v1/rdf/ontology should return 200, got " + response.statusCode()); + String body = response.body(); + assertTrue(body.contains("1.1.0"), "Ontology document should declare the bumped version 1.1.0"); + assertTrue( + body.contains("om:Column") && body.contains("om:TableConstraint"), + "Ontology document should declare core om:Column and om:TableConstraint classes"); + } + + @Test + void testShaclValidateEndpointReturnsReport(TestNamespace ns) throws Exception { + DatabaseService service = DatabaseServiceTestFactory.createPostgres(ns); + DatabaseSchema schema = DatabaseSchemaTestFactory.createSimple(ns, service); + + CreateTable createRequest = new CreateTable(); + createRequest.setName(ns.prefix("rdfShaclTable")); + createRequest.setDatabaseSchema(schema.getFullyQualifiedName()); + createRequest.setColumns(List.of(ColumnBuilder.of("id", "BIGINT").primaryKey().build())); + Table table = Tables.create(createRequest); + + String entityUri = BASE_URI + "entity/table/" + table.getId(); + String url = + SdkClients.getServerUrl() + + "/v1/rdf/validate?entityUri=" + + URLEncoder.encode(entityUri, StandardCharsets.UTF_8); + + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(url)) + .header("Authorization", "Bearer " + SdkClients.getAdminToken()) + .header("Accept", "text/turtle") + .POST(HttpRequest.BodyPublishers.noBody()) + .build(); + + Awaitility.await() + .atMost(Duration.ofSeconds(15)) + .pollInterval(Duration.ofSeconds(1)) + .untilAsserted( + () -> { + HttpResponse response = + HTTP_CLIENT.send(request, HttpResponse.BodyHandlers.ofString()); + assertTrue( + response.statusCode() == 200, + "POST /v1/rdf/validate should return 200, got " + response.statusCode()); + String body = response.body(); + assertTrue( + body.contains("ValidationReport") || body.contains("conforms"), + "Response should be a SHACL validation report"); + assertTrue( + response.headers().firstValue("OM-SHACL-Conforms").isPresent(), + "Endpoint should set OM-SHACL-Conforms header"); + }); + } + + // ---- helpers for the fidelity tests ---- + + private boolean columnIsTypedAsOmColumn(String columnFqn) { + String columnUri = + BASE_URI + "entity/column/" + URLEncoder.encode(columnFqn, StandardCharsets.UTF_8); + String sparql = + String.format( + "PREFIX om: <%s> " + + "ASK { GRAPH ?g { <%s> a om:Column ; om:fullyQualifiedName \"%s\" } }", + OM_NS, columnUri, columnFqn); + return RdfTestUtils.executeSparqlAsk(sparql); + } + + private boolean columnHasBooleanProperty(String columnFqn, String predicate, boolean expected) { + String columnUri = + BASE_URI + "entity/column/" + URLEncoder.encode(columnFqn, StandardCharsets.UTF_8); + String sparql = + String.format( + "PREFIX om: <%s> " + + "PREFIX xsd: " + + "ASK { GRAPH ?g { <%s> om:%s \"%s\"^^xsd:boolean } }", + OM_NS, columnUri, predicate, expected); + return RdfTestUtils.executeSparqlAsk(sparql); + } + + private boolean columnReferencesColumn(String fromFqn, String toFqn) { + String fromUri = + BASE_URI + "entity/column/" + URLEncoder.encode(fromFqn, StandardCharsets.UTF_8); + String toUri = BASE_URI + "entity/column/" + URLEncoder.encode(toFqn, StandardCharsets.UTF_8); + String sparql = + String.format( + "PREFIX om: <%s> ASK { GRAPH ?g { <%s> om:references <%s> } }", OM_NS, fromUri, toUri); + return RdfTestUtils.executeSparqlAsk(sparql); + } + + private boolean tableHasConstraintOfType(String tableFqn, String constraintType) { + String escaped = tableFqn.replace("\\", "\\\\").replace("\"", "\\\""); + String sparql = + String.format( + "PREFIX om: <%s> " + + "ASK { GRAPH ?g { " + + " ?table om:fullyQualifiedName \"%s\" ; " + + " om:hasConstraint ?c . " + + " ?c om:constraintType \"%s\" . " + + "} }", + OM_NS, escaped, constraintType); + return RdfTestUtils.executeSparqlAsk(sparql); + } } diff --git a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/SearchIndexPromotionIT.java b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/SearchIndexPromotionIT.java deleted file mode 100644 index f2bf8a6267dc..000000000000 --- a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/SearchIndexPromotionIT.java +++ /dev/null @@ -1,258 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.it.tests; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.junit.jupiter.api.Assumptions.assumeFalse; - -import java.time.Duration; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Set; -import org.awaitility.Awaitility; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.junit.jupiter.api.parallel.Execution; -import org.junit.jupiter.api.parallel.ExecutionMode; -import org.junit.jupiter.api.parallel.Isolated; -import org.openmetadata.it.bootstrap.TestSuiteBootstrap; -import org.openmetadata.it.factories.DatabaseSchemaTestFactory; -import org.openmetadata.it.factories.DatabaseServiceTestFactory; -import org.openmetadata.it.factories.TableTestFactory; -import org.openmetadata.it.util.SdkClients; -import org.openmetadata.it.util.TestNamespace; -import org.openmetadata.it.util.TestNamespaceExtension; -import org.openmetadata.schema.entity.app.AppRunRecord; -import org.openmetadata.schema.entity.data.DatabaseSchema; -import org.openmetadata.schema.entity.data.Table; -import org.openmetadata.schema.entity.services.DatabaseService; -import org.openmetadata.sdk.network.HttpClient; -import org.openmetadata.sdk.network.HttpMethod; -import org.openmetadata.service.Entity; -import org.openmetadata.service.search.SearchClient; - -@Execution(ExecutionMode.SAME_THREAD) -@Isolated -@ExtendWith(TestNamespaceExtension.class) -public class SearchIndexPromotionIT { - - private static final String APP_NAME = "SearchIndexingApplication"; - private static final String TABLE_ENTITY = "table"; - private static final String TABLE_CANONICAL_ALIAS = "openmetadata_table_search_index"; - private static final String TABLE_SHORT_ALIAS = "openmetadata_table"; - private static final String TABLE_REBUILD_PREFIX = TABLE_CANONICAL_ALIAS + "_rebuild_"; - private static final Set SUCCESS_STATUSES = Set.of("success", "completed"); - private static final Set TERMINAL_STATUSES = - Set.of("success", "completed", "failed", "activeerror", "stopped"); - - @BeforeAll - static void setup() { - SdkClients.adminClient(); - } - - @Test - void tableOnlyRerunPromotesNewStagedIndex(TestNamespace ns) { - assumeFalse( - TestSuiteBootstrap.isK8sEnabled(), "App trigger not compatible with K8s pipeline backend"); - - createTableForReindex(ns); - - HttpClient httpClient = SdkClients.adminClient().getHttpClient(); - waitForCurrentRunCompletion(httpClient); - - String initialTarget = readSingleTableAliasTargetIfPresent(); - Long previousRunStartTime = readLatestRunStartTime(httpClient); - triggerTableReindex(httpClient); - AppRunRecord firstRun = waitForLatestRunSuccess(httpClient, previousRunStartTime); - String firstTarget = waitForPromotedTableAlias(initialTarget); - - triggerTableReindex(httpClient); - waitForLatestRunSuccess(httpClient, firstRun.getStartTime()); - String secondTarget = waitForPromotedTableAlias(firstTarget); - - assertNotEquals(firstTarget, secondTarget, "Second reindex should promote a new staged index"); - assertPreviousTargetIsNotServing(firstTarget); - } - - private static void createTableForReindex(TestNamespace ns) { - DatabaseService service = DatabaseServiceTestFactory.createPostgres(ns); - DatabaseSchema schema = DatabaseSchemaTestFactory.createSimple(ns, service); - Table table = - TableTestFactory.createWithName(ns, schema.getFullyQualifiedName(), "promotion_table"); - - assertNotNull(table.getId(), "Test table should be created before reindex"); - } - - private static void triggerTableReindex(HttpClient httpClient) { - Map config = new HashMap<>(); - config.put("entities", List.of(TABLE_ENTITY)); - config.put("batchSize", 100); - - Awaitility.await("Trigger table-only " + APP_NAME) - .atMost(Duration.ofMinutes(2)) - .pollInterval(Duration.ofSeconds(3)) - .ignoreExceptionsMatching( - e -> e.getMessage() != null && e.getMessage().contains("already running")) - .until( - () -> { - httpClient.execute( - HttpMethod.POST, "/v1/apps/trigger/" + APP_NAME, config, Void.class); - return true; - }); - } - - private static AppRunRecord waitForLatestRunSuccess( - HttpClient httpClient, Long previousRunStartTime) { - AppRunRecord[] holder = new AppRunRecord[1]; - - Awaitility.await("Table reindex run completion") - .atMost(Duration.ofMinutes(5)) - .pollDelay(Duration.ofSeconds(2)) - .pollInterval(Duration.ofSeconds(5)) - .ignoreExceptions() - .untilAsserted( - () -> { - AppRunRecord run = readLatestRun(httpClient); - assertNotNull(run); - assertNotNull(run.getStatus()); - if (previousRunStartTime != null - && run.getStartTime() != null - && run.getStartTime() <= previousRunStartTime) { - throw new AssertionError( - "Latest run is still the pre-trigger one (startTime=" - + run.getStartTime() - + ", previous=" - + previousRunStartTime - + ")"); - } - String status = normalizedStatus(run); - assertTrue( - TERMINAL_STATUSES.contains(status), "Run not in terminal state: " + status); - holder[0] = run; - }); - - AppRunRecord run = holder[0]; - assertTrue( - SUCCESS_STATUSES.contains(normalizedStatus(run)), - () -> "Expected successful table reindex run but got: " + run); - return run; - } - - private static String waitForPromotedTableAlias(String previousTarget) { - String[] target = new String[1]; - - Awaitility.await("Table alias promotion") - .atMost(Duration.ofMinutes(2)) - .pollDelay(Duration.ofSeconds(1)) - .pollInterval(Duration.ofSeconds(2)) - .ignoreExceptions() - .untilAsserted( - () -> { - String currentTarget = readSingleTableAliasTarget(); - assertTrue( - currentTarget.startsWith(TABLE_REBUILD_PREFIX), - "Table alias should point at a staged rebuild index, got " + currentTarget); - if (previousTarget != null) { - assertNotEquals( - previousTarget, - currentTarget, - "Table alias should move to a new staged index after rerun"); - } - Set shortAliasTargets = searchClient().getIndicesByAlias(TABLE_SHORT_ALIAS); - assertTrue( - shortAliasTargets.contains(currentTarget), - "Short table alias should include the promoted staged table index"); - target[0] = currentTarget; - }); - - return target[0]; - } - - private static String readSingleTableAliasTargetIfPresent() { - Set targets = searchClient().getIndicesByAlias(TABLE_CANONICAL_ALIAS); - if (targets.isEmpty()) { - return null; - } - assertEquals(1, targets.size(), "Table canonical alias should have a single target"); - return targets.iterator().next(); - } - - private static String readSingleTableAliasTarget() { - String target = readSingleTableAliasTargetIfPresent(); - assertNotNull(target, "Table canonical alias should point at a promoted index"); - return target; - } - - private static void assertPreviousTargetIsNotServing(String previousTarget) { - SearchClient client = searchClient(); - if (!client.indexExists(previousTarget)) { - return; - } - - Set aliases = client.getAliases(previousTarget); - assertFalse( - aliases.contains(TABLE_CANONICAL_ALIAS), - "Previous staged index should no longer have the canonical table alias"); - assertFalse( - aliases.contains(TABLE_SHORT_ALIAS), - "Previous staged index should no longer have the short table alias"); - } - - private static Long readLatestRunStartTime(HttpClient httpClient) { - try { - AppRunRecord latest = readLatestRun(httpClient); - return latest == null ? null : latest.getStartTime(); - } catch (Exception ignored) { - return null; - } - } - - private static AppRunRecord readLatestRun(HttpClient httpClient) { - return httpClient.execute( - HttpMethod.GET, "/v1/apps/name/" + APP_NAME + "/runs/latest", null, AppRunRecord.class); - } - - private static void waitForCurrentRunCompletion(HttpClient httpClient) { - try { - Awaitility.await("Wait for in-flight " + APP_NAME) - .atMost(Duration.ofMinutes(5)) - .pollInterval(Duration.ofSeconds(3)) - .ignoreExceptions() - .until( - () -> { - AppRunRecord latest = readLatestRun(httpClient); - if (latest == null || latest.getStatus() == null) { - return true; - } - String status = normalizedStatus(latest); - return !"running".equals(status) && !"started".equals(status); - }); - } catch (org.awaitility.core.ConditionTimeoutException ignored) { - // The trigger retry loop handles "already running" if the current run continues. - } - } - - private static String normalizedStatus(AppRunRecord run) { - return run.getStatus().value().toLowerCase(); - } - - private static SearchClient searchClient() { - return Entity.getSearchRepository().getSearchClient(); - } -} diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/McpServer.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/McpServer.java index 9d8a83b4e6c8..9a6711715485 100644 --- a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/McpServer.java +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/McpServer.java @@ -12,14 +12,12 @@ import org.openmetadata.mcp.server.auth.jobs.OAuthTokenCleanupScheduler; import org.openmetadata.mcp.server.transport.OAuthHttpStatelessServerTransportProvider; import org.openmetadata.mcp.tools.DefaultToolContext; -import org.openmetadata.mcp.usage.McpUsageRecorder; import org.openmetadata.schema.utils.JsonUtils; import org.openmetadata.service.Entity; import org.openmetadata.service.OpenMetadataApplicationConfig; import org.openmetadata.service.apps.AbstractNativeApplication; import org.openmetadata.service.apps.ApplicationContext; import org.openmetadata.service.apps.McpServerProvider; -import org.openmetadata.service.apps.bundles.mcp.McpAppConstants; import org.openmetadata.service.limits.Limits; import org.openmetadata.service.security.Authorizer; import org.openmetadata.service.security.ImpersonationContext; @@ -29,7 +27,8 @@ @Slf4j public class McpServer implements McpServerProvider { - private static final String DEFAULT_MCP_BOT_NAME = McpAppConstants.MCP_APP_NAME + "Bot"; + private static final String MCP_APP_NAME = "McpApplication"; + private static final String DEFAULT_MCP_BOT_NAME = MCP_APP_NAME + "Bot"; protected JwtFilter jwtFilter; protected Authorizer authorizer; @@ -222,7 +221,7 @@ private String getMcpBotName() { if (mcpBotName == null) { try { AbstractNativeApplication mcpApp = - ApplicationContext.getInstance().getAppIfExists(McpAppConstants.MCP_APP_NAME); + ApplicationContext.getInstance().getAppIfExists(MCP_APP_NAME); if (mcpApp != null && mcpApp.getApp().getBot() != null) { mcpBotName = mcpApp.getApp().getBot().getName(); } @@ -238,17 +237,12 @@ protected McpStatelessServerFeatures.SyncToolSpecification getTool(McpSchema.Too return new McpStatelessServerFeatures.SyncToolSpecification( tool, (context, req) -> { - CatalogSecurityContext securityContext = - jwtFilter.getCatalogSecurityContext((String) context.get("Authorization")); - String userName = securityContext.getUserPrincipal().getName(); - McpSchema.CallToolResult result = null; try { + CatalogSecurityContext securityContext = + jwtFilter.getCatalogSecurityContext((String) context.get("Authorization")); ImpersonationContext.setImpersonatedBy(getMcpBotName()); - result = toolContext.callTool(authorizer, limits, tool.name(), securityContext, req); - return result; + return toolContext.callTool(authorizer, limits, tool.name(), securityContext, req); } finally { - boolean success = result != null && !Boolean.TRUE.equals(result.isError()); - McpUsageRecorder.record(tool.name(), userName, success); ImpersonationContext.clear(); } }); diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/DefaultToolContext.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/DefaultToolContext.java index ad9d5d2448dd..e75dabe632b4 100644 --- a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/DefaultToolContext.java +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/DefaultToolContext.java @@ -83,6 +83,21 @@ public McpSchema.CallToolResult callTool( case "create_metric": result = new CreateMetricTool().execute(authorizer, limits, securityContext, params); break; + case "sparql_query": + result = new SparqlQueryTool().execute(authorizer, securityContext, params); + break; + case "entity_neighborhood": + result = new EntityNeighborhoodTool().execute(authorizer, securityContext, params); + break; + case "find_by_tag": + result = new FindByTagTool().execute(authorizer, securityContext, params); + break; + case "shacl_validate": + result = new ShaclValidateTool().execute(authorizer, securityContext, params); + break; + case "ontology_describe": + result = new OntologyDescribeTool().execute(authorizer, securityContext, params); + break; default: return McpSchema.CallToolResult.builder() .content( diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/EntityNeighborhoodTool.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/EntityNeighborhoodTool.java new file mode 100644 index 000000000000..43c72ab8c6a2 --- /dev/null +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/EntityNeighborhoodTool.java @@ -0,0 +1,211 @@ +package org.openmetadata.mcp.tools; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.limits.Limits; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +/** + * Returns the n-hop neighborhood of an entity in the knowledge graph as triples. + * + *

Driven by a SPARQL CONSTRUCT with bounded property paths so depth is enforced even for + * adversarial inputs. Depth is hard-capped at 3 (matches REST graph-explorer). + */ +@Slf4j +public class EntityNeighborhoodTool implements McpTool { + + private static final int MIN_DEPTH = 1; + private static final int MAX_DEPTH = 3; + private static final int DEFAULT_DEPTH = 2; + + @Override + public Map execute( + Authorizer authorizer, CatalogSecurityContext securityContext, Map params) + throws IOException { + String entityId = string(params, "entityId"); + String entityType = string(params, "entityType"); + if (entityId == null || entityId.isBlank()) { + return error("'entityId' parameter is required"); + } + if (entityType == null || entityType.isBlank()) { + return error("'entityType' parameter is required"); + } + try { + UUID.fromString(entityId); + } catch (IllegalArgumentException e) { + return error("'entityId' must be a UUID"); + } + if (!entityType.matches("[A-Za-z][A-Za-z0-9]*")) { + return error("'entityType' must be alphanumeric"); + } + + int depth = clampDepth(intParam(params, "depth", DEFAULT_DEPTH)); + int limit = Math.min(Math.max(intParam(params, "limit", 200), 1), 2000); + + RdfRepository repository = RdfRepository.getInstanceOrNull(); + if (repository == null || !repository.isEnabled()) { + return error("RDF repository is not enabled on this OpenMetadata server"); + } + + String entityUri = repository.getBaseUri() + "entity/" + entityType + "/" + entityId; + + String constructQuery = buildConstructQuery(entityUri, depth, limit); + String triples; + try { + triples = repository.executeSparqlQuery(constructQuery, "text/turtle"); + } catch (Exception e) { + LOG.error("CONSTRUCT for neighborhood failed for {}", entityUri, e); + return error("Neighborhood query failed: " + e.getMessage()); + } + + String selectQuery = buildSelectQuery(entityUri, depth, limit); + String selectJson; + try { + selectJson = repository.executeSparqlQuery(selectQuery, "application/sparql-results+json"); + } catch (Exception e) { + LOG.error("SELECT for neighborhood failed for {}", entityUri, e); + selectJson = null; + } + + Map result = new LinkedHashMap<>(); + result.put("entityUri", entityUri); + result.put("depth", depth); + result.put("limit", limit); + result.put("triples", triples == null ? "" : triples); + result.put("edges", parseEdges(selectJson)); + return result; + } + + @Override + public Map execute( + Authorizer authorizer, + Limits limits, + CatalogSecurityContext securityContext, + Map params) { + throw new UnsupportedOperationException( + "EntityNeighborhoodTool does not enforce write limits."); + } + + /** + * Builds a CONSTRUCT query that yields all triples on every path of length 1..depth radiating + * from the entity. Each UNION arm emits exactly one triple template ({@code ?s ?p ?o}), with + * {@code ?s} bound to the *actual* subject of that triple — not the start entity. The previous + * implementation bound {@code ?s = } unconditionally, which collapsed all 2- and + * 3-hop edges onto the start node and produced an incorrect graph. + */ + static String buildConstructQuery(String entityUri, int depth, int limit) { + String e = "<" + entityUri + ">"; + StringBuilder w = new StringBuilder(); + // depth-1 outgoing + w.append(" { BIND(").append(e).append(" AS ?s) ").append(e).append(" ?p ?o }\n"); + // depth-1 incoming + w.append(" UNION { BIND(").append(e).append(" AS ?o) ?s ?p ").append(e).append(" }\n"); + if (depth >= 2) { + // depth-2 outgoing: emit BOTH edges as separate arms + w.append(" UNION { BIND(") + .append(e) + .append(" AS ?s) ") + .append(e) + .append(" ?p ?o . ?o ?p2 ?n2 }\n"); + w.append(" UNION { ").append(e).append(" ?p1 ?s . ?s ?p ?o }\n"); + } + if (depth >= 3) { + // depth-3 outgoing: three edges, three arms + w.append(" UNION { BIND(") + .append(e) + .append(" AS ?s) ") + .append(e) + .append(" ?p ?o . ?o ?p2 ?n2 . ?n2 ?p3 ?n3 }\n"); + w.append(" UNION { ").append(e).append(" ?p1 ?s . ?s ?p ?o . ?o ?p3 ?n3 }\n"); + w.append(" UNION { ").append(e).append(" ?p1 ?n1 . ?n1 ?p2 ?s . ?s ?p ?o }\n"); + } + return "CONSTRUCT { ?s ?p ?o } WHERE {\n" + w + "} LIMIT " + limit; + } + + static String buildSelectQuery(String entityUri, int depth, int limit) { + return "PREFIX om: \n" + + "SELECT ?direction ?predicate ?neighbor ?neighborLabel WHERE {\n" + + " { BIND('outgoing' AS ?direction) <" + + entityUri + + "> ?predicate ?neighbor }\n" + + " UNION { BIND('incoming' AS ?direction) ?neighbor ?predicate <" + + entityUri + + "> }\n" + + " OPTIONAL { ?neighbor ?neighborLabel }\n" + + "} LIMIT " + + limit; + } + + private static List> parseEdges(String selectJson) { + if (selectJson == null || selectJson.isBlank()) { + return List.of(); + } + try { + Map sparql = JsonUtils.readValue(selectJson, Map.class); + Object results = sparql.get("results"); + if (!(results instanceof Map resultsMap)) return List.of(); + Object bindings = resultsMap.get("bindings"); + if (!(bindings instanceof List rows)) return List.of(); + List> edges = new ArrayList<>(rows.size()); + for (Object row : rows) { + if (!(row instanceof Map r)) continue; + Map edge = new LinkedHashMap<>(); + edge.put("direction", bindingValue(r, "direction")); + edge.put("predicate", bindingValue(r, "predicate")); + edge.put("neighbor", bindingValue(r, "neighbor")); + Object label = bindingValue(r, "neighborLabel"); + if (label != null) { + edge.put("neighborLabel", label); + } + edges.add(edge); + } + return edges; + } catch (Exception e) { + LOG.warn("Failed to parse neighborhood SELECT results: {}", e.getMessage()); + return List.of(); + } + } + + private static Object bindingValue(Map row, String name) { + Object node = row.get(name); + if (!(node instanceof Map nodeMap)) return null; + return nodeMap.get("value"); + } + + private static int clampDepth(int depth) { + return Math.min(Math.max(depth, MIN_DEPTH), MAX_DEPTH); + } + + private static int intParam(Map params, String key, int defaultValue) { + Object v = params.get(key); + if (v instanceof Number n) return n.intValue(); + if (v instanceof String s) { + try { + return Integer.parseInt(s); + } catch (NumberFormatException e) { + return defaultValue; + } + } + return defaultValue; + } + + private static String string(Map params, String key) { + Object v = params.get(key); + return v instanceof String s ? s : null; + } + + private static Map error(String message) { + Map result = new HashMap<>(); + result.put("error", message); + return result; + } +} diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/FindByTagTool.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/FindByTagTool.java new file mode 100644 index 000000000000..824f70544139 --- /dev/null +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/FindByTagTool.java @@ -0,0 +1,168 @@ +package org.openmetadata.mcp.tools; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.limits.Limits; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +/** + * Finds entities tagged with a given tag or glossary term FQN. + * + *

Resolves the tag/glossary URI from the FQN, then runs a SELECT for everything connected + * via {@code om:hasTag} or {@code om:hasGlossaryTerm}. Results include the entity URI, its FQN, + * its rdf:type, and its label. + */ +@Slf4j +public class FindByTagTool implements McpTool { + + @Override + public Map execute( + Authorizer authorizer, CatalogSecurityContext securityContext, Map params) + throws IOException { + String tagFqn = string(params, "tagFqn"); + if (tagFqn == null || tagFqn.isBlank()) { + return error("'tagFqn' parameter is required"); + } + if (tagFqn.contains("\"") || tagFqn.contains("\\") || tagFqn.contains("\n")) { + return error("'tagFqn' contains illegal characters"); + } + int limit = Math.min(Math.max(intParam(params, "limit", 50), 1), 500); + int offset = Math.max(intParam(params, "offset", 0), 0); + String entityType = string(params, "entityType"); + if (entityType != null && !entityType.matches("[A-Za-z][A-Za-z0-9]*")) { + return error("'entityType' must be alphanumeric"); + } + + RdfRepository repository = RdfRepository.getInstanceOrNull(); + if (repository == null || !repository.isEnabled()) { + return error("RDF repository is not enabled on this OpenMetadata server"); + } + + String sparql = buildSparql(tagFqn, entityType, limit, offset); + String json; + try { + json = repository.executeSparqlQuery(sparql, "application/sparql-results+json"); + } catch (Exception e) { + LOG.error("find_by_tag SPARQL failed for {}", tagFqn, e); + return error("SPARQL execution failed: " + e.getMessage()); + } + + List> entities = parseRows(json); + Map result = new LinkedHashMap<>(); + result.put("tagFqn", tagFqn); + result.put("entityTypeFilter", entityType); + result.put("limit", limit); + result.put("offset", offset); + result.put("results", entities); + result.put("returnedCount", entities.size()); + return result; + } + + @Override + public Map execute( + Authorizer authorizer, + Limits limits, + CatalogSecurityContext securityContext, + Map params) { + throw new UnsupportedOperationException("FindByTagTool does not enforce write limits."); + } + + static String buildSparql(String tagFqn, String entityType, int limit, int offset) { + String escapedFqn = tagFqn.replace("\"", "\\\""); + StringBuilder sb = new StringBuilder(); + // Match either a Tag (om:tagFQN) or a GlossaryTerm (om:fullyQualifiedName) — the input FQN + // can be either. GlossaryTerms in OM RDF do not carry om:tagFQN, so without this UNION the + // tool silently returned zero results for any glossary FQN. + sb.append("PREFIX om: \n") + .append("PREFIX rdfs: \n") + .append("SELECT DISTINCT ?entity ?entityType ?fqn ?label WHERE {\n") + .append(" { ?tag om:tagFQN \"") + .append(escapedFqn) + .append("\" }\n") + .append(" UNION { ?tag om:fullyQualifiedName \"") + .append(escapedFqn) + .append("\" }\n") + .append(" { ?entity om:hasTag ?tag } UNION { ?entity om:hasGlossaryTerm ?tag }\n") + .append(" ?entity a ?entityType .\n") + .append(" OPTIONAL { ?entity om:fullyQualifiedName ?fqn }\n") + .append(" OPTIONAL { ?entity rdfs:label ?label }\n"); + if (entityType != null && !entityType.isBlank()) { + String capitalized = Character.toUpperCase(entityType.charAt(0)) + entityType.substring(1); + sb.append(" FILTER(?entityType = )\n"); + } + sb.append("} ORDER BY ?fqn LIMIT ").append(limit).append(" OFFSET ").append(offset); + return sb.toString(); + } + + private static List> parseRows(String selectJson) { + if (selectJson == null || selectJson.isBlank()) { + return List.of(); + } + try { + Map sparql = JsonUtils.readValue(selectJson, Map.class); + Object results = sparql.get("results"); + if (!(results instanceof Map resultsMap)) return List.of(); + Object bindings = resultsMap.get("bindings"); + if (!(bindings instanceof List rows)) return List.of(); + List> entities = new ArrayList<>(rows.size()); + for (Object row : rows) { + if (!(row instanceof Map r)) continue; + Map entity = new LinkedHashMap<>(); + Object e = bindingValue(r, "entity"); + if (e == null) continue; + entity.put("entity", e); + Object t = bindingValue(r, "entityType"); + if (t != null) entity.put("entityType", t); + Object fqn = bindingValue(r, "fqn"); + if (fqn != null) entity.put("fullyQualifiedName", fqn); + Object label = bindingValue(r, "label"); + if (label != null) entity.put("label", label); + entities.add(entity); + } + return entities; + } catch (Exception e) { + LOG.warn("Failed to parse find_by_tag SELECT results: {}", e.getMessage()); + return List.of(); + } + } + + private static Object bindingValue(Map row, String name) { + Object node = row.get(name); + if (!(node instanceof Map nodeMap)) return null; + return nodeMap.get("value"); + } + + private static String string(Map params, String key) { + Object v = params.get(key); + return v instanceof String s ? s : null; + } + + private static int intParam(Map params, String key, int defaultValue) { + Object v = params.get(key); + if (v instanceof Number n) return n.intValue(); + if (v instanceof String s) { + try { + return Integer.parseInt(s); + } catch (NumberFormatException e) { + return defaultValue; + } + } + return defaultValue; + } + + private static Map error(String message) { + Map result = new HashMap<>(); + result.put("error", message); + return result; + } +} diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/OntologyDescribeTool.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/OntologyDescribeTool.java new file mode 100644 index 000000000000..302b91f9fa17 --- /dev/null +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/OntologyDescribeTool.java @@ -0,0 +1,116 @@ +package org.openmetadata.mcp.tools; + +import java.io.IOException; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.service.limits.Limits; +import org.openmetadata.service.rdf.RdfIriValidator; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.resources.rdf.OntologyDocument; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +/** + * Returns either the entire OpenMetadata ontology or a focused DESCRIBE for a single class / + * property URI. + * + *

The full ontology is served from the bundled classpath resource (no triplestore round + * trip). A focused DESCRIBE goes through the triplestore so it picks up any side-ontology + * extensions registered there. + */ +@Slf4j +public class OntologyDescribeTool implements McpTool { + + @Override + public Map execute( + Authorizer authorizer, CatalogSecurityContext securityContext, Map params) + throws IOException { + String resource = string(params, "resource"); + String format = normalizeFormat(string(params, "format")); + + if (resource == null || resource.isBlank()) { + OntologyDocument.SerializedOntology serialized = OntologyDocument.serializeAsString(format); + Map result = new LinkedHashMap<>(); + result.put("scope", "full-ontology"); + result.put("format", format); + result.put("mediaType", serialized.mediaType()); + result.put("body", serialized.body()); + return result; + } + + String validatedResource = RdfIriValidator.validateEntityIri(resource); + if (validatedResource == null) { + return error( + "'resource' must be a valid absolute http(s) IRI (no whitespace, control characters," + + " angle brackets, or quotes)"); + } + + RdfRepository repository = RdfRepository.getInstanceOrNull(); + if (repository == null || !repository.isEnabled()) { + return error("RDF repository is not enabled; cannot DESCRIBE individual ontology resources"); + } + + String describe = "DESCRIBE <" + validatedResource + ">"; + String mime = formatMime(format); + String body; + try { + body = repository.executeSparqlQueryDirect(describe, mime); + } catch (Exception e) { + LOG.error("Ontology DESCRIBE failed for {}", validatedResource, e); + return error("DESCRIBE failed: " + e.getMessage()); + } + + Map result = new LinkedHashMap<>(); + result.put("scope", "describe"); + result.put("resource", validatedResource); + result.put("format", format); + result.put("mediaType", mime); + result.put("body", body == null ? "" : body); + return result; + } + + @Override + public Map execute( + Authorizer authorizer, + Limits limits, + CatalogSecurityContext securityContext, + Map params) { + throw new UnsupportedOperationException("OntologyDescribeTool does not enforce write limits."); + } + + private static String normalizeFormat(String format) { + if (format == null) return "turtle"; + return switch (format.toLowerCase()) { + case "jsonld", "json-ld", "ld+json" -> "jsonld"; + case "rdfxml", "rdf+xml", "rdf/xml" -> "rdfxml"; + case "ntriples", "n-triples" -> "ntriples"; + default -> "turtle"; + }; + } + + /** + * Maps a normalised format name to the SPARQL-accept MIME type the triplestore expects. Every + * format returned by {@link #normalizeFormat} must round-trip through here. + */ + private static String formatMime(String format) { + return switch (format) { + case "jsonld" -> "application/ld+json"; + case "rdfxml" -> "application/rdf+xml"; + case "ntriples" -> "application/n-triples"; + default -> "text/turtle"; + }; + } + + private static String string(Map params, String key) { + Object v = params.get(key); + return v instanceof String s ? s : null; + } + + private static Map error(String message) { + Map result = new HashMap<>(); + result.put("error", message); + return result; + } +} diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/ShaclValidateTool.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/ShaclValidateTool.java new file mode 100644 index 000000000000..e33421f4233a --- /dev/null +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/ShaclValidateTool.java @@ -0,0 +1,153 @@ +package org.openmetadata.mcp.tools; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.StringReader; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.UUID; +import lombok.extern.slf4j.Slf4j; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.ModelFactory; +import org.apache.jena.riot.Lang; +import org.apache.jena.riot.RDFDataMgr; +import org.apache.jena.riot.RDFFormat; +import org.apache.jena.shacl.ValidationReport; +import org.openmetadata.service.limits.Limits; +import org.openmetadata.service.rdf.RdfIriValidator; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.resources.rdf.RdfShaclValidator; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +/** + * Runs SHACL validation against either a single entity's subgraph or the entire dataset and + * returns the resulting validation report. + * + *

If {@code entityId} + {@code entityType} (or {@code entityUri}) is supplied, the tool runs + * {@code DESCRIBE } first and validates only that subgraph. Otherwise it pulls everything + * and validates the whole graph (admin-style usage; expensive). + */ +@Slf4j +public class ShaclValidateTool implements McpTool { + + @Override + public Map execute( + Authorizer authorizer, CatalogSecurityContext securityContext, Map params) + throws IOException { + RdfRepository repository = RdfRepository.getInstanceOrNull(); + if (repository == null || !repository.isEnabled()) { + return error("RDF repository is not enabled on this OpenMetadata server"); + } + + String entityUri = resolveEntityUri(params, repository.getBaseUri()); + if (entityUri != null && entityUri.startsWith("error:")) { + return error(entityUri.substring("error:".length())); + } + + if (entityUri == null && !Boolean.TRUE.equals(params.get("fullGraph"))) { + return error( + "Full-graph SHACL validation must be explicitly enabled by passing fullGraph=true. " + + "It loads the entire triplestore into memory and can OOM the MCP server. Prefer " + + "passing entityId+entityType (or entityUri) to scope the check."); + } + + String constructQuery = + entityUri != null + ? "DESCRIBE <" + entityUri + ">" + : "CONSTRUCT { ?s ?p ?o } WHERE { ?s ?p ?o }"; + + String dataTurtle; + try { + dataTurtle = repository.executeSparqlQueryDirect(constructQuery, "text/turtle"); + } catch (Exception e) { + LOG.error("Subgraph fetch for SHACL failed", e); + return error("Failed to fetch subgraph: " + e.getMessage()); + } + + Model dataModel = ModelFactory.createDefaultModel(); + try (StringReader reader = new StringReader(dataTurtle == null ? "" : dataTurtle)) { + RDFDataMgr.read(dataModel, reader, repository.getBaseUri(), Lang.TURTLE); + } catch (Exception e) { + LOG.error("Failed to parse subgraph for SHACL validation", e); + return error("Failed to parse subgraph: " + e.getMessage()); + } + + ValidationReport report = RdfShaclValidator.validate(dataModel); + String format = normalizeFormat(string(params, "format")); + RDFFormat rdfFormat = + "jsonld".equals(format) ? RDFFormat.JSONLD_PRETTY : RDFFormat.TURTLE_PRETTY; + ByteArrayOutputStream out = new ByteArrayOutputStream(); + RDFDataMgr.write(out, report.getModel(), rdfFormat); + + Map result = new LinkedHashMap<>(); + result.put("scope", entityUri == null ? "full-graph" : "entity"); + if (entityUri != null) { + result.put("entityUri", entityUri); + } + result.put("conforms", report.conforms()); + long violationCount = report.getEntries() == null ? 0 : report.getEntries().stream().count(); + result.put("violationCount", violationCount); + result.put("format", format); + result.put("report", out.toString(StandardCharsets.UTF_8)); + return result; + } + + @Override + public Map execute( + Authorizer authorizer, + Limits limits, + CatalogSecurityContext securityContext, + Map params) { + throw new UnsupportedOperationException("ShaclValidateTool does not enforce write limits."); + } + + private static String resolveEntityUri(Map params, String baseUri) { + String entityUri = string(params, "entityUri"); + if (entityUri != null && !entityUri.isBlank()) { + String validated = RdfIriValidator.validateEntityIri(entityUri); + if (validated == null) { + return "error:'entityUri' must be a valid absolute http(s) IRI"; + } + return validated; + } + String entityId = string(params, "entityId"); + String entityType = string(params, "entityType"); + if (entityId == null && entityType == null) { + return null; // full-graph scope + } + if (entityId == null || entityType == null) { + return "error:Both 'entityId' and 'entityType' are required when scoping by entity, or omit both for full-graph scope"; + } + try { + UUID.fromString(entityId); + } catch (IllegalArgumentException e) { + return "error:'entityId' must be a UUID"; + } + if (!entityType.matches("[A-Za-z][A-Za-z0-9]*")) { + return "error:'entityType' must be alphanumeric"; + } + return baseUri + "entity/" + entityType + "/" + entityId; + } + + private static String normalizeFormat(String format) { + if (format == null) return "turtle"; + return switch (format.toLowerCase()) { + case "jsonld", "json-ld", "ld+json" -> "jsonld"; + default -> "turtle"; + }; + } + + private static String string(Map params, String key) { + Object v = params.get(key); + return v instanceof String s ? s : null; + } + + private static Map error(String message) { + Map result = new HashMap<>(); + result.put("error", message); + return result; + } +} diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/SparqlQueryTool.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/SparqlQueryTool.java new file mode 100644 index 000000000000..f1dfaf0bd8c4 --- /dev/null +++ b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/tools/SparqlQueryTool.java @@ -0,0 +1,181 @@ +package org.openmetadata.mcp.tools; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import lombok.extern.slf4j.Slf4j; +import org.apache.jena.query.Query; +import org.apache.jena.query.QueryException; +import org.apache.jena.query.QueryFactory; +import org.openmetadata.service.limits.Limits; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.rdf.federation.SparqlFederationGuard; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +/** + * Read-only SPARQL query tool for AI agents. Wraps {@code RdfRepository.executeSparqlQuery}. + * + *

This tool is deliberately the only path through which an MCP client can talk SPARQL. + * Hardening: + * + *

+ */ +@Slf4j +public class SparqlQueryTool implements McpTool { + + private static final int DEFAULT_MAX_BYTES = 1 * 1024 * 1024; + private static final int HARD_MAX_BYTES = 16 * 1024 * 1024; + + @Override + public Map execute( + Authorizer authorizer, CatalogSecurityContext securityContext, Map params) + throws IOException { + String query = string(params, "query"); + if (query == null || query.isBlank()) { + return error("'query' parameter is required"); + } + + Query parsed; + try { + parsed = QueryFactory.create(query); + } catch (QueryException e) { + return error("SPARQL parse error: " + e.getMessage()); + } + if (!isReadOnly(parsed)) { + return error( + "Only read-only SPARQL queries (SELECT, ASK, DESCRIBE, CONSTRUCT) are allowed via this tool. Use the admin REST endpoint for SPARQL UPDATE."); + } + + RdfRepository repository = RdfRepository.getInstanceOrNull(); + if (repository == null || !repository.isEnabled()) { + return error("RDF repository is not enabled on this OpenMetadata server"); + } + + try { + new SparqlFederationGuard(repository.getConfig()).enforce(query); + } catch (SparqlFederationGuard.FederationDisallowedException e) { + return error(e.getMessage()); + } + + String format = normalizeFormat(string(params, "format")); + String mimeType = mimeFor(format); + String inferenceLevel = string(params, "inferenceLevel"); + + int maxBytes = clamp(intParam(params, "maxBytes", DEFAULT_MAX_BYTES), 1024, HARD_MAX_BYTES); + + String body; + try { + body = + inferenceLevel != null + && !inferenceLevel.isBlank() + && !"none".equalsIgnoreCase(inferenceLevel) + ? repository.executeSparqlQueryWithInference(query, mimeType, inferenceLevel) + : repository.executeSparqlQuery(query, mimeType); + } catch (Exception e) { + LOG.error("SPARQL query execution failed", e); + return error("SPARQL execution failed: " + e.getMessage()); + } + + Map result = new HashMap<>(); + result.put("format", format); + result.put("queryType", parsed.queryType().toString()); + if (body == null) { + result.put("body", ""); + result.put("truncated", false); + result.put("byteCount", 0); + return result; + } + byte[] bytes = body.getBytes(java.nio.charset.StandardCharsets.UTF_8); + boolean truncated = bytes.length > maxBytes; + if (truncated) { + // Truncate by bytes (not chars). Multi-byte UTF-8 sequences must not be split mid-rune, + // so back off until we land on the start of a code point (top bits != 10xxxxxx). + int cut = maxBytes; + while (cut > 0 && (bytes[cut] & 0xC0) == 0x80) { + cut--; + } + result.put("body", new String(bytes, 0, cut, java.nio.charset.StandardCharsets.UTF_8)); + } else { + result.put("body", body); + } + result.put("truncated", truncated); + result.put("byteCount", bytes.length); + return result; + } + + @Override + public Map execute( + Authorizer authorizer, + Limits limits, + CatalogSecurityContext securityContext, + Map params) { + throw new UnsupportedOperationException("SparqlQueryTool does not enforce write limits."); + } + + private static boolean isReadOnly(Query query) { + return query.isSelectType() + || query.isAskType() + || query.isDescribeType() + || query.isConstructType(); + } + + private static String normalizeFormat(String format) { + if (format == null || format.isBlank()) { + return "json"; + } + return switch (format.toLowerCase()) { + case "json", "xml", "csv", "tsv", "turtle", "rdfxml", "ntriples", "jsonld" -> format + .toLowerCase(); + default -> "json"; + }; + } + + private static String mimeFor(String format) { + return switch (format) { + case "xml" -> "application/sparql-results+xml"; + case "csv" -> "text/csv"; + case "tsv" -> "text/tab-separated-values"; + case "turtle" -> "text/turtle"; + case "rdfxml" -> "application/rdf+xml"; + case "ntriples" -> "application/n-triples"; + case "jsonld" -> "application/ld+json"; + default -> "application/sparql-results+json"; + }; + } + + private static String string(Map params, String key) { + Object v = params.get(key); + return v instanceof String s ? s : null; + } + + private static int intParam(Map params, String key, int defaultValue) { + Object v = params.get(key); + if (v instanceof Number n) return n.intValue(); + if (v instanceof String s) { + try { + return Integer.parseInt(s); + } catch (NumberFormatException e) { + return defaultValue; + } + } + return defaultValue; + } + + private static int clamp(int v, int lo, int hi) { + return Math.min(Math.max(v, lo), hi); + } + + private static Map error(String message) { + Map result = new HashMap<>(); + result.put("error", message); + return result; + } +} diff --git a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/usage/McpUsageRecorder.java b/openmetadata-mcp/src/main/java/org/openmetadata/mcp/usage/McpUsageRecorder.java deleted file mode 100644 index 94fdbdccb48d..000000000000 --- a/openmetadata-mcp/src/main/java/org/openmetadata/mcp/usage/McpUsageRecorder.java +++ /dev/null @@ -1,79 +0,0 @@ -/* - * Copyright 2025 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.mcp.usage; - -import org.openmetadata.schema.entity.app.App; -import org.openmetadata.schema.entity.app.AppExtension; -import org.openmetadata.schema.entity.app.mcp.McpToolCallUsage; -import org.openmetadata.schema.utils.JsonUtils; -import org.openmetadata.service.Entity; -import org.openmetadata.service.apps.AbstractNativeApplication; -import org.openmetadata.service.apps.ApplicationContext; -import org.openmetadata.service.apps.bundles.mcp.McpAppConstants; -import org.openmetadata.service.jdbi3.CollectionDAO; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Best-effort one-row-per-call writer for MCP tool invocations. Records to the - * {@code apps_extension_time_series} table reusing the {@code limits} extension type — the same - * per-app usage bucket CollateAI writes to. Rows are isolated from other apps by - * {@code appName='McpApplication'}, so the shared extension causes no cross-talk. Pure tracking. - * No billing, no enforcement, no rate-limiting. A recording failure must never break the tool - * call, so every code path catches and logs. - */ -public final class McpUsageRecorder { - - private static final Logger LOG = LoggerFactory.getLogger(McpUsageRecorder.class); - - private McpUsageRecorder() {} - - public static void record(String toolName, String userName, boolean success) { - try { - App app = resolveMcpApp(); - if (app == null) { - LOG.debug( - "McpApplication not initialized, skipping MCP usage record for tool {}", toolName); - return; - } - McpToolCallUsage usage = - new McpToolCallUsage() - .withAppId(app.getId()) - .withAppName(app.getName()) - .withTimestamp(System.currentTimeMillis()) - .withExtension(AppExtension.ExtensionType.LIMITS) - .withToolName(toolName) - .withUserName(userName) - .withSuccess(success); - getDao().insert(JsonUtils.pojoToJson(usage), AppExtension.ExtensionType.LIMITS.toString()); - } catch (Exception e) { - LOG.warn( - "Failed to record MCP usage for tool={} user={} success={}: {}", - toolName, - userName, - success, - e.getMessage()); - } - } - - private static App resolveMcpApp() { - AbstractNativeApplication app = - ApplicationContext.getInstance().getAppIfExists(McpAppConstants.MCP_APP_NAME); - return app != null ? app.getApp() : null; - } - - private static CollectionDAO.AppExtensionTimeSeries getDao() { - return Entity.getCollectionDAO().appExtensionTimeSeriesDao(); - } -} diff --git a/openmetadata-mcp/src/main/resources/json/data/mcp/tools.json b/openmetadata-mcp/src/main/resources/json/data/mcp/tools.json index d88b67d37ac8..658fa5b38d7e 100644 --- a/openmetadata-mcp/src/main/resources/json/data/mcp/tools.json +++ b/openmetadata-mcp/src/main/resources/json/data/mcp/tools.json @@ -584,6 +584,135 @@ "entityType" ] } + }, + { + "name": "sparql_query", + "description": "Run a read-only SPARQL query (SELECT, ASK, DESCRIBE, CONSTRUCT) against the OpenMetadata knowledge graph. Use this when the question is graph-shaped — e.g. 'all tables with FK references into customers.id', 'all upstream tables of dashboard X', 'columns tagged PII whose null count > 0'. UPDATE / INSERT / DELETE / DROP / LOAD / CLEAR / CREATE are rejected; use the admin REST endpoint for writes. SERVICE clauses against external endpoints are rejected unless on the federation allowlist. Returns the result body in the chosen format (default 'json' for SELECT/ASK; 'turtle' is recommended for DESCRIBE/CONSTRUCT). Body is capped at maxBytes (default 1 MiB).", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The SPARQL query string. Use https://open-metadata.org/ontology/ as the om: prefix." + }, + "format": { + "type": "string", + "description": "Result format. SELECT/ASK: json (default), xml, csv, tsv. CONSTRUCT/DESCRIBE: turtle, jsonld, ntriples, rdfxml.", + "default": "json" + }, + "inferenceLevel": { + "type": "string", + "description": "Optional. Reasoning level applied at query time: none (default), rdfs, owl, custom.", + "default": "none" + }, + "maxBytes": { + "type": "integer", + "description": "Cap on the response body size in bytes. Default 1048576 (1 MiB). Min 1024, max 16777216.", + "default": 1048576 + } + }, + "required": ["query"] + } + }, + { + "name": "entity_neighborhood", + "description": "Return the n-hop neighborhood of an OpenMetadata entity in the knowledge graph as triples plus a flat edges list. Use this when you have a specific entity (UUID + type) and want to see what else is connected to it: hasColumn, belongsToSchema, hasTag, owners, lineage, etc. Depth is bounded at 1-3 hops. Returns Turtle for the triples and a structured edges array (direction, predicate, neighbor URI, neighbor label).", + "parameters": { + "type": "object", + "properties": { + "entityId": { + "type": "string", + "description": "UUID of the entity (the 'id' field on the entity)." + }, + "entityType": { + "type": "string", + "description": "Entity type singular (table, dashboard, pipeline, glossaryTerm, ...)." + }, + "depth": { + "type": "integer", + "description": "Hop depth, 1-3. Default 2.", + "default": 2 + }, + "limit": { + "type": "integer", + "description": "Maximum triples returned. Default 200, max 2000.", + "default": 200 + } + }, + "required": ["entityId", "entityType"] + } + }, + { + "name": "find_by_tag", + "description": "Find all entities tagged with a given classification tag or glossary term FQN. Walks om:hasTag and om:hasGlossaryTerm. Use this for 'show me everything classified PII.Sensitive' or 'all assets associated with the BusinessTerms.Customer glossary term'. Optionally filter by entity type. Paginated.", + "parameters": { + "type": "object", + "properties": { + "tagFqn": { + "type": "string", + "description": "FQN of the tag or glossary term (e.g. 'PII.Sensitive', 'BusinessTerms.Customer')." + }, + "entityType": { + "type": "string", + "description": "Optional. Restrict results to a single entity type (e.g. 'table')." + }, + "limit": { + "type": "integer", + "description": "Maximum results. Default 50, max 500.", + "default": 50 + }, + "offset": { + "type": "integer", + "description": "Page offset. Default 0.", + "default": 0 + } + }, + "required": ["tagFqn"] + } + }, + { + "name": "shacl_validate", + "description": "Run SHACL validation against either a single entity's subgraph (scoped via entityId+entityType or entityUri) or the entire dataset. Returns conforms (boolean), violationCount, and the SHACL ValidationReport in the requested RDF format. Read-only; never blocks writes. Use this to debug data-quality issues before mutating the graph.", + "parameters": { + "type": "object", + "properties": { + "entityId": { + "type": "string", + "description": "Optional. UUID of the entity to scope validation to." + }, + "entityType": { + "type": "string", + "description": "Optional. Entity type singular. Required if entityId is provided." + }, + "entityUri": { + "type": "string", + "description": "Optional alternative to entityId+entityType: full URI of the entity." + }, + "format": { + "type": "string", + "description": "Report format: turtle (default) or jsonld.", + "default": "turtle" + } + } + } + }, + { + "name": "ontology_describe", + "description": "Return the OpenMetadata ontology — either the full canonical ontology (when 'resource' is omitted) or a SPARQL DESCRIBE for a single class or property URI. Use this to discover available classes, properties, ranges, and domains the agent can query. Format: turtle (default), jsonld, ntriples, rdfxml.", + "parameters": { + "type": "object", + "properties": { + "resource": { + "type": "string", + "description": "Optional. Full URI of the class or property to DESCRIBE (e.g. 'https://open-metadata.org/ontology/Column'). Omit to return the full ontology." + }, + "format": { + "type": "string", + "description": "Output format. Default 'turtle'.", + "default": "turtle" + } + } + } } ] } diff --git a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/EntityNeighborhoodToolTest.java b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/EntityNeighborhoodToolTest.java new file mode 100644 index 000000000000..6aa44c932632 --- /dev/null +++ b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/EntityNeighborhoodToolTest.java @@ -0,0 +1,179 @@ +package org.openmetadata.mcp.tools; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Map; +import java.util.UUID; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +class EntityNeighborhoodToolTest { + + private static final Authorizer AUTHORIZER = mock(Authorizer.class); + private static final CatalogSecurityContext SEC = mock(CatalogSecurityContext.class); + + @Test + @DisplayName("Missing entityId is rejected") + void missingEntityId() throws IOException { + Map result = + new EntityNeighborhoodTool().execute(AUTHORIZER, SEC, Map.of("entityType", "table")); + assertEquals("'entityId' parameter is required", result.get("error")); + } + + @Test + @DisplayName("Missing entityType is rejected") + void missingEntityType() throws IOException { + Map result = + new EntityNeighborhoodTool() + .execute(AUTHORIZER, SEC, Map.of("entityId", UUID.randomUUID().toString())); + assertEquals("'entityType' parameter is required", result.get("error")); + } + + @Test + @DisplayName("Non-UUID entityId is rejected with a clean error") + void nonUuidEntityIdRejected() throws IOException { + Map result = + new EntityNeighborhoodTool() + .execute(AUTHORIZER, SEC, Map.of("entityId", "not-a-uuid", "entityType", "table")); + assertEquals("'entityId' must be a UUID", result.get("error")); + } + + @Test + @DisplayName("Non-alphanumeric entityType is rejected (defends against URI injection)") + void badEntityTypeRejected() throws IOException { + Map result = + new EntityNeighborhoodTool() + .execute( + AUTHORIZER, + SEC, + Map.of("entityId", UUID.randomUUID().toString(), "entityType", "table> ; DROP --")); + assertEquals("'entityType' must be alphanumeric", result.get("error")); + } + + @Test + @DisplayName("RDF disabled returns service-unavailable error") + void rdfDisabled() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(false); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new EntityNeighborhoodTool() + .execute( + AUTHORIZER, + SEC, + Map.of("entityId", UUID.randomUUID().toString(), "entityType", "table")); + assertNotNull(result.get("error")); + } + } + + @Test + @DisplayName("Successful call returns triples + edges + clamped depth") + void successfulCall() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + when(repo.executeSparqlQuery(anyString(), org.mockito.ArgumentMatchers.eq("text/turtle"))) + .thenReturn( + " ."); + when(repo.executeSparqlQuery( + anyString(), org.mockito.ArgumentMatchers.eq("application/sparql-results+json"))) + .thenReturn( + "{\"results\":{\"bindings\":[" + + "{\"direction\":{\"value\":\"outgoing\"},\"predicate\":{\"value\":\"https://open-metadata.org/ontology/hasColumn\"},\"neighbor\":{\"value\":\"urn:c1\"}}" + + "]}}"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new EntityNeighborhoodTool() + .execute( + AUTHORIZER, + SEC, + Map.of( + "entityId", + "11111111-1111-1111-1111-111111111111", + "entityType", + "table", + "depth", + 99)); + + assertNull(result.get("error")); + assertEquals(3, result.get("depth"), "Depth must be clamped to 3"); + assertNotNull(result.get("triples")); + @SuppressWarnings("unchecked") + var edges = (java.util.List>) result.get("edges"); + assertEquals(1, edges.size()); + assertEquals("outgoing", edges.get(0).get("direction")); + } + } + + @Test + @DisplayName("buildConstructQuery includes inverse direction (incoming edges)") + void constructQueryIncludesInverse() { + String q = + EntityNeighborhoodTool.buildConstructQuery( + "https://open-metadata.org/entity/table/abc", 2, 100); + // Incoming arm now binds ?o = and matches ?s ?p + assertTrue(q.contains("?s ?p ")); + } + + @Test + @DisplayName("buildConstructQuery respects depth") + void constructQueryDepthBounds() { + String d1 = EntityNeighborhoodTool.buildConstructQuery("urn:e", 1, 100); + String d3 = EntityNeighborhoodTool.buildConstructQuery("urn:e", 3, 100); + // depth-1: only outgoing + incoming arms (split by UNION = 2 pieces + 1 = 3) + assertTrue(d1.split("UNION").length <= 3, "Depth 1 should not contain 2/3-hop unions"); + assertTrue(d3.contains("?n3"), "Depth 3 must include the 3-hop chain variable"); + } + + @Test + @DisplayName("buildConstructQuery preserves real subjects on multi-hop arms") + void constructQueryPreservesMultiHopSubjects() { + String q = EntityNeighborhoodTool.buildConstructQuery("urn:e", 2, 100); + // The 2-hop "second-edge" arm must NOT bind ?s to the start entity — it must let the + // intermediate node be ?s so the emitted triple is faithful to the real graph. + assertTrue( + q.contains(" ?p1 ?s . ?s ?p ?o"), + "Depth-2 second-edge arm must bind ?s to the intermediate node, not "); + } + + @Test + @DisplayName("Repository throws → tool returns clean error") + void repositoryThrows() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenThrow(new RuntimeException("boom")); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new EntityNeighborhoodTool() + .execute( + AUTHORIZER, + SEC, + Map.of( + "entityId", + "22222222-2222-2222-2222-222222222222", + "entityType", + "pipeline")); + assertTrue(((String) result.get("error")).contains("Neighborhood query failed")); + } + } +} diff --git a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/FindByTagToolTest.java b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/FindByTagToolTest.java new file mode 100644 index 000000000000..544da8d611a7 --- /dev/null +++ b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/FindByTagToolTest.java @@ -0,0 +1,159 @@ +package org.openmetadata.mcp.tools; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Map; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +class FindByTagToolTest { + + private static final Authorizer AUTHORIZER = mock(Authorizer.class); + private static final CatalogSecurityContext SEC = mock(CatalogSecurityContext.class); + + @Test + @DisplayName("Missing tagFqn is rejected") + void missingTagFqn() throws IOException { + Map result = new FindByTagTool().execute(AUTHORIZER, SEC, Map.of()); + assertEquals("'tagFqn' parameter is required", result.get("error")); + } + + @Test + @DisplayName("tagFqn containing a quote is rejected (defends against SPARQL string injection)") + void illegalQuoteRejected() throws IOException { + Map result = + new FindByTagTool().execute(AUTHORIZER, SEC, Map.of("tagFqn", "PII\".Sensitive")); + assertEquals("'tagFqn' contains illegal characters", result.get("error")); + } + + @Test + @DisplayName("tagFqn with backslash or newline is rejected") + void illegalControlCharsRejected() throws IOException { + Map result = + new FindByTagTool().execute(AUTHORIZER, SEC, Map.of("tagFqn", "PII.Sens\nitive")); + assertEquals("'tagFqn' contains illegal characters", result.get("error")); + + Map result2 = + new FindByTagTool().execute(AUTHORIZER, SEC, Map.of("tagFqn", "PII\\Sensitive")); + assertEquals("'tagFqn' contains illegal characters", result2.get("error")); + } + + @Test + @DisplayName("Non-alphanumeric entityType is rejected") + void badEntityTypeRejected() throws IOException { + Map result = + new FindByTagTool() + .execute( + AUTHORIZER, SEC, Map.of("tagFqn", "PII.Sensitive", "entityType", "table OR 1=1")); + assertEquals("'entityType' must be alphanumeric", result.get("error")); + } + + @Test + @DisplayName("RDF disabled returns service-unavailable error") + void rdfDisabled() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(false); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new FindByTagTool().execute(AUTHORIZER, SEC, Map.of("tagFqn", "PII.Sensitive")); + assertNotNull(result.get("error")); + } + } + + @Test + @DisplayName("Successful call returns parsed entities with FQN, type, label") + void successfulCall() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenReturn( + "{\"results\":{\"bindings\":[" + + "{\"entity\":{\"value\":\"https://open-metadata.org/entity/table/abc\"}," + + " \"entityType\":{\"value\":\"https://open-metadata.org/ontology/Table\"}," + + " \"fqn\":{\"value\":\"svc.db.s.t\"}," + + " \"label\":{\"value\":\"t\"}}" + + "]}}"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new FindByTagTool().execute(AUTHORIZER, SEC, Map.of("tagFqn", "PII.Sensitive")); + assertNull(result.get("error")); + assertEquals(1, result.get("returnedCount")); + @SuppressWarnings("unchecked") + java.util.List> rows = + (java.util.List>) result.get("results"); + assertEquals("svc.db.s.t", rows.get(0).get("fullyQualifiedName")); + } + } + + @Test + @DisplayName( + "buildSparql escapes embedded quotes via the regex earlier; the final query is parameterized correctly") + void buildSparqlContainsEscapedFqn() { + String sparql = FindByTagTool.buildSparql("PII.Sensitive", "table", 50, 0); + assertTrue(sparql.contains("\"PII.Sensitive\"")); + assertTrue(sparql.contains("LIMIT 50")); + assertTrue(sparql.contains("OFFSET 0")); + assertTrue(sparql.contains("ontology/Table")); + } + + @Test + @DisplayName("buildSparql matches GlossaryTerm by om:fullyQualifiedName, not only om:tagFQN") + void buildSparqlMatchesGlossaryFqn() { + String sparql = FindByTagTool.buildSparql("MyGlossary.PII", null, 50, 0); + assertTrue( + sparql.contains("om:tagFQN \"MyGlossary.PII\""), "must still match Tags by om:tagFQN"); + assertTrue( + sparql.contains("om:fullyQualifiedName \"MyGlossary.PII\""), + "must also match GlossaryTerms by om:fullyQualifiedName"); + } + + @Test + @DisplayName("Empty result set returns empty list, not error") + void emptyResultSet() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenReturn("{\"results\":{\"bindings\":[]}}"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new FindByTagTool().execute(AUTHORIZER, SEC, Map.of("tagFqn", "PII.None")); + assertNull(result.get("error")); + assertEquals(0, result.get("returnedCount")); + } + } + + @Test + @DisplayName("Limit beyond hard max is clamped to 500") + void limitClamped() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenReturn("{\"results\":{\"bindings\":[]}}"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new FindByTagTool() + .execute(AUTHORIZER, SEC, Map.of("tagFqn", "PII.None", "limit", 999_999)); + assertEquals(500, result.get("limit")); + } + } +} diff --git a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/OntologyDescribeToolTest.java b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/OntologyDescribeToolTest.java new file mode 100644 index 000000000000..d8e303f587dc --- /dev/null +++ b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/OntologyDescribeToolTest.java @@ -0,0 +1,93 @@ +package org.openmetadata.mcp.tools; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Map; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +class OntologyDescribeToolTest { + + private static final Authorizer AUTHORIZER = mock(Authorizer.class); + private static final CatalogSecurityContext SEC = mock(CatalogSecurityContext.class); + + @Test + @DisplayName( + "No 'resource' returns the full ontology from classpath without touching the triplestore") + void fullOntologyServedFromClasspath() throws IOException { + Map result = new OntologyDescribeTool().execute(AUTHORIZER, SEC, Map.of()); + assertNull(result.get("error")); + assertEquals("full-ontology", result.get("scope")); + String body = (String) result.get("body"); + assertNotNull(body); + assertTrue( + body.contains("om:") || body.contains("ontology"), + "Full ontology body must look like RDF, got: " + + body.substring(0, Math.min(200, body.length()))); + } + + @Test + @DisplayName("Non-URI 'resource' is rejected") + void nonUriResourceRejected() throws IOException { + Map result = + new OntologyDescribeTool().execute(AUTHORIZER, SEC, Map.of("resource", "Column")); + assertEquals( + "'resource' must be an absolute http(s) URI for the class or property", + result.get("error")); + } + + @Test + @DisplayName("RDF disabled while DESCRIBE-ing a single class returns service-unavailable error") + void describeRequiresRdf() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(null); + Map result = + new OntologyDescribeTool() + .execute( + AUTHORIZER, SEC, Map.of("resource", "https://open-metadata.org/ontology/Column")); + assertNotNull(result.get("error")); + } + } + + @Test + @DisplayName("Successful DESCRIBE call returns turtle by default") + void successfulDescribe() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.executeSparqlQueryDirect( + anyString(), org.mockito.ArgumentMatchers.eq("text/turtle"))) + .thenReturn("@prefix om: ."); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new OntologyDescribeTool() + .execute( + AUTHORIZER, SEC, Map.of("resource", "https://open-metadata.org/ontology/Column")); + assertEquals("describe", result.get("scope")); + assertEquals("turtle", result.get("format")); + assertEquals("text/turtle", result.get("mediaType")); + assertNotNull(result.get("body")); + } + } + + @Test + @DisplayName("Format normalizes 'json-ld' to 'jsonld'") + void formatNormalization() throws IOException { + Map result = + new OntologyDescribeTool().execute(AUTHORIZER, SEC, Map.of("format", "json-ld")); + assertEquals("jsonld", result.get("format")); + } +} diff --git a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/ShaclValidateToolTest.java b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/ShaclValidateToolTest.java new file mode 100644 index 000000000000..ab9e98e7bd83 --- /dev/null +++ b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/ShaclValidateToolTest.java @@ -0,0 +1,179 @@ +package org.openmetadata.mcp.tools; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Map; +import java.util.UUID; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +class ShaclValidateToolTest { + + private static final Authorizer AUTHORIZER = mock(Authorizer.class); + private static final CatalogSecurityContext SEC = mock(CatalogSecurityContext.class); + + @Test + @DisplayName("RDF disabled returns service-unavailable error") + void rdfDisabled() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(null); + Map result = new ShaclValidateTool().execute(AUTHORIZER, SEC, Map.of()); + assertNotNull(result.get("error")); + } + } + + @Test + @DisplayName("Non-URI entityUri is rejected") + void nonUriEntityUri() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new ShaclValidateTool().execute(AUTHORIZER, SEC, Map.of("entityUri", "not-an-http-uri")); + assertEquals("'entityUri' must be an absolute http(s) URI", result.get("error")); + } + } + + @Test + @DisplayName("entityId without entityType is rejected") + void entityIdWithoutType() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new ShaclValidateTool() + .execute(AUTHORIZER, SEC, Map.of("entityId", UUID.randomUUID().toString())); + assertNotNull(result.get("error")); + } + } + + @Test + @DisplayName("Non-UUID entityId is rejected") + void badEntityId() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new ShaclValidateTool() + .execute(AUTHORIZER, SEC, Map.of("entityId", "abc", "entityType", "table")); + assertEquals("'entityId' must be a UUID", result.get("error")); + } + } + + @Test + @DisplayName("Successful entity-scoped validation reports conforms with violationCount") + void successfulEntityScopedValidation() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + // Conforming subgraph: a typed entity with the required label and FQN. + String turtleSubgraph = + "@prefix om: .\n" + + "@prefix rdfs: .\n" + + " a om:Table ;\n" + + " rdfs:label \"abc\" ;\n" + + " om:fullyQualifiedName \"svc.db.s.abc\" ;\n" + + " om:hasColumn .\n" + + " a om:Column ;\n" + + " om:fullyQualifiedName \"svc.db.s.abc.id\" ."; + when(repo.executeSparqlQueryDirect( + anyString(), org.mockito.ArgumentMatchers.eq("text/turtle"))) + .thenReturn(turtleSubgraph); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new ShaclValidateTool() + .execute( + AUTHORIZER, + SEC, + Map.of( + "entityId", "11111111-1111-1111-1111-111111111111", "entityType", "table")); + assertNull(result.get("error")); + assertEquals("entity", result.get("scope")); + assertNotNull(result.get("conforms")); + assertNotNull(result.get("report")); + } + } + + @Test + @DisplayName("Subgraph that violates a shape returns conforms=false and a non-empty report") + void violationDetected() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + // Bad column lineage: fromColumn is a literal where shape requires om:Column. + String bad = + "@prefix om: .\n" + + " a om:ColumnLineage ;\n" + + " om:fromColumn \"svc.db.s.t.col_a\" ."; + when(repo.executeSparqlQueryDirect(anyString(), anyString())).thenReturn(bad); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new ShaclValidateTool() + .execute(AUTHORIZER, SEC, Map.of("entityUri", "https://open-metadata.org/lin/1")); + assertEquals(false, result.get("conforms")); + assertTrue(((Long) result.get("violationCount")) > 0); + } + } + + @Test + @DisplayName("Empty body from triplestore is handled gracefully") + void emptyBodyHandled() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + when(repo.executeSparqlQueryDirect(anyString(), anyString())).thenReturn(null); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + new ShaclValidateTool() + .execute( + AUTHORIZER, + SEC, + Map.of( + "entityId", "11111111-1111-1111-1111-111111111111", "entityType", "table")); + assertNull(result.get("error")); + assertNotNull(result.get("conforms")); + } + } + + @Test + @DisplayName("No scope params → full-graph validation") + void fullGraphScope() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getBaseUri()).thenReturn("https://open-metadata.org/"); + when(repo.executeSparqlQueryDirect(anyString(), anyString())).thenReturn(""); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = new ShaclValidateTool().execute(AUTHORIZER, SEC, Map.of()); + assertEquals("full-graph", result.get("scope")); + } + } +} diff --git a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/SparqlQueryToolTest.java b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/SparqlQueryToolTest.java new file mode 100644 index 000000000000..2941d37eeb2b --- /dev/null +++ b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/tools/SparqlQueryToolTest.java @@ -0,0 +1,256 @@ +package org.openmetadata.mcp.tools; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Map; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.openmetadata.schema.api.configuration.rdf.RdfConfiguration; +import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.security.Authorizer; +import org.openmetadata.service.security.auth.CatalogSecurityContext; + +/** + * Failure-mode tests for SparqlQueryTool. Each test names the bad input it is exercising — + * these are the queries we expect adversarial or sloppy MCP clients to send. + */ +class SparqlQueryToolTest { + + private static final Authorizer AUTHORIZER = mock(Authorizer.class); + private static final CatalogSecurityContext SEC = mock(CatalogSecurityContext.class); + + private static SparqlQueryTool newTool() { + return new SparqlQueryTool(); + } + + @Test + @DisplayName("Missing 'query' parameter returns a clean error") + void missingQueryParam() throws IOException { + Map result = newTool().execute(AUTHORIZER, SEC, Map.of()); + assertEquals("'query' parameter is required", result.get("error")); + } + + @Test + @DisplayName("Empty / blank query is rejected") + void blankQueryRejected() throws IOException { + Map result = newTool().execute(AUTHORIZER, SEC, Map.of("query", " ")); + assertNotNull(result.get("error")); + } + + @Test + @DisplayName("Garbage SPARQL surfaces as a parse error rather than a 500") + void garbageQuerySurfacesParseError() throws IOException { + Map result = + newTool().execute(AUTHORIZER, SEC, Map.of("query", "not sparql at all {{}}")); + assertNotNull(result.get("error")); + assertTrue(((String) result.get("error")).startsWith("SPARQL parse error")); + } + + @Test + @DisplayName("INSERT DATA is rejected — only read-only queries allowed via this tool") + void insertDataRejected() throws IOException { + String q = "INSERT DATA { }"; + Map result = newTool().execute(AUTHORIZER, SEC, Map.of("query", q)); + String err = (String) result.get("error"); + // INSERT DATA fails the SPARQL Query parser (it's an update operation), so we surface a + // parse error. Either way, the tool refuses it. + assertNotNull(err); + } + + @Test + @DisplayName("DELETE WHERE is rejected — only read-only queries allowed via this tool") + void deleteRejected() throws IOException { + String q = "DELETE WHERE { ?s ?p ?o }"; + Map result = newTool().execute(AUTHORIZER, SEC, Map.of("query", q)); + assertNotNull(result.get("error")); + } + + @Test + @DisplayName("DROP GRAPH is rejected") + void dropRejected() throws IOException { + String q = "DROP GRAPH "; + Map result = newTool().execute(AUTHORIZER, SEC, Map.of("query", q)); + assertNotNull(result.get("error")); + } + + @Test + @DisplayName("RDF disabled on the server returns a service-unavailable error") + void repositoryDisabled() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(false); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + newTool().execute(AUTHORIZER, SEC, Map.of("query", "SELECT * WHERE { ?s ?p ?o }")); + assertEquals( + "RDF repository is not enabled on this OpenMetadata server", result.get("error")); + } + } + + @Test + @DisplayName("RDF instance missing returns a service-unavailable error") + void repositoryMissing() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(null); + + Map result = + newTool().execute(AUTHORIZER, SEC, Map.of("query", "SELECT * WHERE { ?s ?p ?o }")); + assertNotNull(result.get("error")); + } + } + + @Test + @DisplayName("SERVICE clause to non-allowlisted endpoint is rejected") + void serviceClauseRejected() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getConfig()).thenReturn(new RdfConfiguration()); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + String q = "SELECT * WHERE { SERVICE { ?s ?p ?o } }"; + Map result = newTool().execute(AUTHORIZER, SEC, Map.of("query", q)); + assertNotNull(result.get("error")); + assertTrue(((String) result.get("error")).contains("SERVICE")); + } + } + + @Test + @DisplayName("Successful SELECT returns body, format, queryType, and untruncated metadata") + void successfulSelect() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getConfig()).thenReturn(new RdfConfiguration()); + when(repo.executeSparqlQuery( + "SELECT * WHERE { ?s ?p ?o } LIMIT 1", "application/sparql-results+json")) + .thenReturn("{\"results\":{\"bindings\":[]}}"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + newTool() + .execute(AUTHORIZER, SEC, Map.of("query", "SELECT * WHERE { ?s ?p ?o } LIMIT 1")); + + assertNull(result.get("error")); + assertEquals("json", result.get("format")); + assertEquals("SELECT", result.get("queryType")); + assertFalse((Boolean) result.get("truncated")); + assertNotNull(result.get("body")); + } + } + + @Test + @DisplayName("Result larger than maxBytes is truncated and flagged") + void resultTruncatedWhenOversized() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getConfig()).thenReturn(new RdfConfiguration()); + String hugeBody = "x".repeat(8_000); + when(repo.executeSparqlQuery( + "SELECT * WHERE { ?s ?p ?o }", "application/sparql-results+json")) + .thenReturn(hugeBody); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + newTool() + .execute( + AUTHORIZER, + SEC, + Map.of("query", "SELECT * WHERE { ?s ?p ?o }", "maxBytes", 2048)); + assertTrue((Boolean) result.get("truncated")); + assertTrue(((String) result.get("body")).length() <= 2048); + assertEquals(8_000, result.get("byteCount")); + } + } + + @Test + @DisplayName("Repository throws → tool returns clean error rather than propagating") + void repositoryThrows() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getConfig()).thenReturn(new RdfConfiguration()); + when(repo.executeSparqlQuery( + org.mockito.ArgumentMatchers.anyString(), org.mockito.ArgumentMatchers.anyString())) + .thenThrow(new RuntimeException("Fuseki connection refused")); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + newTool().execute(AUTHORIZER, SEC, Map.of("query", "ASK { ?s ?p ?o }")); + String err = (String) result.get("error"); + assertNotNull(err); + assertTrue(err.contains("SPARQL execution failed")); + } + } + + @Test + @DisplayName("Inference level 'rdfs' routes through the inference path") + void inferenceLevelRouted() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getConfig()).thenReturn(new RdfConfiguration()); + when(repo.executeSparqlQueryWithInference( + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.eq("rdfs"))) + .thenReturn("{}"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + newTool() + .execute( + AUTHORIZER, + SEC, + Map.of("query", "SELECT * WHERE { ?s ?p ?o }", "inferenceLevel", "rdfs")); + assertNull(result.get("error")); + org.mockito.Mockito.verify(repo) + .executeSparqlQueryWithInference( + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.eq("rdfs")); + } + } + + @Test + @DisplayName("Format defaults to json when unspecified or unrecognized") + void formatDefaultsToJson() throws IOException { + try (MockedStatic mocked = mockStatic(RdfRepository.class)) { + RdfRepository repo = mock(RdfRepository.class); + when(repo.isEnabled()).thenReturn(true); + when(repo.getConfig()).thenReturn(new RdfConfiguration()); + when(repo.executeSparqlQuery( + org.mockito.ArgumentMatchers.anyString(), + org.mockito.ArgumentMatchers.eq("application/sparql-results+json"))) + .thenReturn("{}"); + mocked.when(RdfRepository::getInstanceOrNull).thenReturn(repo); + + Map result = + newTool() + .execute( + AUTHORIZER, + SEC, + Map.of("query", "SELECT * WHERE { ?s ?p ?o }", "format", "weirdformat")); + assertEquals("json", result.get("format")); + } + } + + @Test + @DisplayName("Limits-aware execute throws — write-tool contract not applicable") + void writeContractNotSupported() { + org.junit.jupiter.api.Assertions.assertThrows( + UnsupportedOperationException.class, + () -> newTool().execute(AUTHORIZER, null, SEC, Map.of("query", "ASK {}"))); + } +} diff --git a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/usage/McpUsageRecorderTest.java b/openmetadata-mcp/src/test/java/org/openmetadata/mcp/usage/McpUsageRecorderTest.java deleted file mode 100644 index 0eb2c01e1add..000000000000 --- a/openmetadata-mcp/src/test/java/org/openmetadata/mcp/usage/McpUsageRecorderTest.java +++ /dev/null @@ -1,153 +0,0 @@ -/* - * Copyright 2025 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.mcp.usage; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.util.UUID; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; -import org.mockito.MockedStatic; -import org.mockito.Mockito; -import org.openmetadata.schema.entity.app.App; -import org.openmetadata.schema.entity.app.AppExtension; -import org.openmetadata.schema.entity.app.mcp.McpToolCallUsage; -import org.openmetadata.schema.utils.JsonUtils; -import org.openmetadata.service.Entity; -import org.openmetadata.service.apps.AbstractNativeApplication; -import org.openmetadata.service.apps.ApplicationContext; -import org.openmetadata.service.apps.bundles.mcp.McpAppConstants; -import org.openmetadata.service.jdbi3.CollectionDAO; - -class McpUsageRecorderTest { - - private CollectionDAO.AppExtensionTimeSeries dao; - private MockedStatic entityStatic; - private MockedStatic appContextStatic; - private ApplicationContext appContext; - - @BeforeEach - void setUp() { - dao = mock(CollectionDAO.AppExtensionTimeSeries.class); - CollectionDAO collectionDAO = mock(CollectionDAO.class); - when(collectionDAO.appExtensionTimeSeriesDao()).thenReturn(dao); - - entityStatic = Mockito.mockStatic(Entity.class); - entityStatic.when(Entity::getCollectionDAO).thenReturn(collectionDAO); - - appContext = mock(ApplicationContext.class); - appContextStatic = Mockito.mockStatic(ApplicationContext.class); - appContextStatic.when(ApplicationContext::getInstance).thenReturn(appContext); - } - - @AfterEach - void tearDown() { - entityStatic.close(); - appContextStatic.close(); - } - - @Test - void recordWritesUsageRowWhenAppRegistered() { - UUID appId = UUID.randomUUID(); - stubMcpApp(appId, McpAppConstants.MCP_APP_NAME); - - long before = System.currentTimeMillis(); - McpUsageRecorder.record("search_metadata", "alice", true); - long after = System.currentTimeMillis(); - - ArgumentCaptor json = ArgumentCaptor.forClass(String.class); - ArgumentCaptor ext = ArgumentCaptor.forClass(String.class); - verify(dao, times(1)).insert(json.capture(), ext.capture()); - assertThat(ext.getValue()).isEqualTo("limits"); - - McpToolCallUsage decoded = JsonUtils.readValue(json.getValue(), McpToolCallUsage.class); - assertThat(decoded.getAppId()).isEqualTo(appId); - assertThat(decoded.getAppName()).isEqualTo(McpAppConstants.MCP_APP_NAME); - assertThat(decoded.getToolName()).isEqualTo("search_metadata"); - assertThat(decoded.getUserName()).isEqualTo("alice"); - assertThat(decoded.getSuccess()).isTrue(); - assertThat(decoded.getExtension()).isEqualTo(AppExtension.ExtensionType.LIMITS); - assertThat(decoded.getTimestamp()).isBetween(before, after); - } - - /** - * The {@code apps_extension_time_series} table has generated columns {@code appId}, - * {@code appName}, and {@code timestamp} that read from the JSON payload using those exact - * property names. If the serialized field names ever drift (rename, missing field) the rows - * still insert but the columns become null, breaking every read query. Lock the on-the-wire - * names so the contract is checked at build time rather than via a failing prod query. - */ - @Test - void serializedJsonContainsGeneratedColumnFieldNames() { - stubMcpApp(UUID.randomUUID(), McpAppConstants.MCP_APP_NAME); - - McpUsageRecorder.record("any_tool", "alice", true); - - ArgumentCaptor json = ArgumentCaptor.forClass(String.class); - verify(dao).insert(json.capture(), eq("limits")); - String raw = json.getValue(); - assertThat(raw).contains("\"appId\":"); - assertThat(raw).contains("\"appName\":"); - assertThat(raw).contains("\"timestamp\":"); - assertThat(raw).contains("\"extension\":\"limits\""); - } - - @Test - void recordSkipsWhenMcpApplicationNotInitialized() { - when(appContext.getAppIfExists(McpAppConstants.MCP_APP_NAME)).thenReturn(null); - - McpUsageRecorder.record("any_tool", "alice", true); - - verify(dao, never()).insert(anyString(), anyString()); - } - - @Test - void recordSwallowsDaoException() { - stubMcpApp(UUID.randomUUID(), McpAppConstants.MCP_APP_NAME); - doThrow(new RuntimeException("db down")).when(dao).insert(anyString(), eq("limits")); - - McpUsageRecorder.record("create_glossary", "alice", false); - - verify(dao, times(1)).insert(anyString(), eq("limits")); - } - - @Test - void recordCapturesFailureFlag() { - stubMcpApp(UUID.randomUUID(), McpAppConstants.MCP_APP_NAME); - - McpUsageRecorder.record("patch_entity", "bob", false); - - ArgumentCaptor json = ArgumentCaptor.forClass(String.class); - verify(dao).insert(json.capture(), eq("limits")); - McpToolCallUsage decoded = JsonUtils.readValue(json.getValue(), McpToolCallUsage.class); - assertThat(decoded.getSuccess()).isFalse(); - } - - private void stubMcpApp(UUID appId, String appName) { - AbstractNativeApplication nativeApp = mock(AbstractNativeApplication.class); - App app = new App().withId(appId).withName(appName); - when(nativeApp.getApp()).thenReturn(app); - when(appContext.getAppIfExists(appName)).thenReturn(nativeApp); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java index 8710ab3f55d1..c7a783078d20 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java @@ -89,6 +89,7 @@ import org.openmetadata.service.apps.scheduler.AppScheduler; import org.openmetadata.service.audit.AuditLogEventPublisher; import org.openmetadata.service.audit.AuditLogRepository; +import org.openmetadata.service.cache.CacheConfig; import org.openmetadata.service.config.CacheConfiguration; import org.openmetadata.service.config.OMWebBundle; import org.openmetadata.service.config.OMWebConfiguration; @@ -396,7 +397,7 @@ public void run(OpenMetadataApplicationConfig catalogConfig, Environment environ jdbi.onDemand(CollectionDAO.class), Entity.getSearchRepository())); // Register Distributed Job Participant for distributed search indexing - registerDistributedJobParticipant(environment, jdbi); + registerDistributedJobParticipant(environment, jdbi, catalogConfig.getCacheConfig()); registerDistributedRdfJobParticipant(environment, jdbi); // Register Event publishers @@ -1131,18 +1132,24 @@ private void initializeWebsockets( } } - protected void registerDistributedJobParticipant(Environment environment, Jdbi jdbi) { + protected void registerDistributedJobParticipant( + Environment environment, Jdbi jdbi, CacheConfig cacheConfig) { try { CollectionDAO collectionDAO = jdbi.onDemand(CollectionDAO.class); SearchRepository searchRepository = Entity.getSearchRepository(); String serverId = ServerIdentityResolver.getInstance().getServerId(); DistributedJobParticipant participant = - new DistributedJobParticipant(collectionDAO, searchRepository, serverId); + new DistributedJobParticipant(collectionDAO, searchRepository, serverId, cacheConfig); environment.lifecycle().manage(participant); + String notifierType = + (cacheConfig != null && cacheConfig.provider == CacheConfig.Provider.redis) + ? "Redis Pub/Sub" + : "database polling"; LOG.info( - "Registered DistributedJobParticipant for distributed search indexing using database polling"); + "Registered DistributedJobParticipant for distributed search indexing using {}", + notifierType); } catch (Exception e) { LOG.warn("Failed to register DistributedJobParticipant", e); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/mcp/McpAppConstants.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/mcp/McpAppConstants.java deleted file mode 100644 index 8952dc9b2720..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/mcp/McpAppConstants.java +++ /dev/null @@ -1,26 +0,0 @@ -/* - * Copyright 2025 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.service.apps.bundles.mcp; - -/** - * Single source of truth for the {@link McpApplication} name written into - * {@code apps_extension_time_series.appName} and queried by the read-side resource. Prevents - * value drift across the recorder, the MCP server, and the REST resource. - */ -public final class McpAppConstants { - - public static final String MCP_APP_NAME = "McpApplication"; - - private McpAppConstants() {} -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoff.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoff.java new file mode 100644 index 000000000000..bac6039873c7 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoff.java @@ -0,0 +1,35 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +/** + * Replaces fixed-delay sleep in backpressure loops with exponential backoff. Starts at an initial + * delay and doubles on each call up to a configurable maximum. Call {@link #reset()} when + * backpressure clears so the next occurrence starts fresh. + */ +public class AdaptiveBackoff { + + private final long initialMs; + private final long maxMs; + private long currentMs; + + public AdaptiveBackoff(long initialMs, long maxMs) { + if (initialMs <= 0) { + throw new IllegalArgumentException("initialMs must be > 0"); + } + if (maxMs < initialMs) { + throw new IllegalArgumentException("maxMs must be >= initialMs"); + } + this.initialMs = initialMs; + this.maxMs = maxMs; + this.currentMs = initialMs; + } + + public long nextDelay() { + long delay = currentMs; + currentMs = Math.min(currentMs * 2, maxMs); + return delay; + } + + public void reset() { + currentMs = initialMs; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java index 83c3d4b8ee04..6fd00d9d1641 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategy.java @@ -1,7 +1,13 @@ package org.openmetadata.service.apps.bundles.searchIndex; +import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; +import static org.openmetadata.service.Entity.QUERY_COST_RECORD; +import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; +import static org.openmetadata.service.Entity.TEST_CASE_RESULT; + import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -14,6 +20,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.analytics.ReportData; import org.openmetadata.schema.system.EventPublisherJob; import org.openmetadata.schema.system.Stats; import org.openmetadata.schema.system.StepStats; @@ -32,7 +39,19 @@ import org.openmetadata.service.util.FullyQualifiedName; @Slf4j -public class DistributedIndexingStrategy { +public class DistributedIndexingStrategy implements IndexingStrategy { + + private static final Set TIME_SERIES_ENTITIES = + Set.of( + ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), + ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), + ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), + TEST_CASE_RESOLUTION_STATUS, + TEST_CASE_RESULT, + QUERY_COST_RECORD); + private static final long MONITOR_POLL_INTERVAL_MS = 2000; private final CollectionDAO collectionDAO; @@ -41,7 +60,6 @@ public class DistributedIndexingStrategy { private final UUID appId; private final Long appStartTime; private final String createdBy; - private final DistributedReindexStatsMapper statsMapper; private final CompositeProgressListener listeners = new CompositeProgressListener(); private final AtomicBoolean stopped = new AtomicBoolean(false); @@ -64,13 +82,14 @@ public DistributedIndexingStrategy( this.appId = appId; this.appStartTime = appStartTime; this.createdBy = createdBy; - this.statsMapper = new DistributedReindexStatsMapper(collectionDAO); } + @Override public void addListener(ReindexingProgressListener listener) { listeners.addListener(listener); } + @Override public ExecutionResult execute(ReindexingConfiguration config, ReindexingJobContext context) { long startTime = System.currentTimeMillis(); try { @@ -94,10 +113,9 @@ private ExecutionResult doExecute( ReindexingConfiguration config, ReindexingJobContext context, long startTime) { this.config = config; - Set entityTypes = SearchIndexEntityTypes.normalizeEntityTypes(config.entities()); - LOG.info("Starting distributed reindexing for entities: {}", entityTypes); + LOG.info("Starting distributed reindexing for entities: {}", config.entities()); - Stats stats = initializeTotalRecords(entityTypes); + Stats stats = initializeTotalRecords(config.entities()); currentStats.set(stats); int partitionSize = jobData.getPartitionSize() != null ? jobData.getPartitionSize() : 10000; @@ -107,7 +125,7 @@ private ExecutionResult doExecute( distributedExecutor.addListener(listeners); SearchIndexJob distributedJob = - distributedExecutor.createJob(entityTypes, jobData, createdBy, config); + distributedExecutor.createJob(config.entities(), jobData, createdBy, config); LOG.info( "Created distributed job {} with {} total records", @@ -118,19 +136,21 @@ private ExecutionResult doExecute( searchRepository.createBulkSink( config.batchSize(), config.maxConcurrentRequests(), config.payloadSize()); - RecreateIndexHandler stagedIndexHandler = searchRepository.createReindexHandler(); - if (stagedIndexHandler instanceof DefaultRecreateHandler defaultHandler) { + RecreateIndexHandler recreateIndexHandler = searchRepository.createReindexHandler(); + if (recreateIndexHandler instanceof DefaultRecreateHandler defaultHandler) { defaultHandler.withJobData(jobData); } - ReindexContext stagedIndexContext = stagedIndexHandler.reCreateIndexes(entityTypes); - if (stagedIndexContext == null || stagedIndexContext.isEmpty()) { - throw new IllegalStateException( - "Staged index preparation did not produce any target indexes"); + ReindexContext recreateContext = null; + + if (config.recreateIndex()) { + recreateContext = recreateIndexHandler.reCreateIndexes(config.entities()); + if (recreateContext != null && !recreateContext.isEmpty()) { + distributedExecutor.updateStagedIndexMapping(recreateContext.getStagedIndexMapping()); + } } - distributedExecutor.updateStagedIndexMapping(stagedIndexContext.getStagedIndexMapping()); distributedExecutor.setAppContext(appId, appStartTime); - distributedExecutor.execute(searchIndexSink, stagedIndexContext, config); + distributedExecutor.execute(searchIndexSink, recreateContext, config.recreateIndex(), config); monitorDistributedJob(distributedJob.getId()); @@ -161,8 +181,8 @@ private ExecutionResult doExecute( boolean success = finalizeAllEntityReindex( - stagedIndexHandler, - stagedIndexContext, + recreateIndexHandler, + recreateContext, !stopped.get() && !hasIncompleteProcessing(stats)); ExecutionResult.Status resultStatus = determineStatus(stats); @@ -270,7 +290,166 @@ private void monitorDistributedJob(UUID jobId) { private void updateStatsFromDistributedJob( Stats stats, SearchIndexJob distributedJob, StepStats actualSinkStats) { - statsMapper.updateStats(stats, distributedJob, actualSinkStats, getColumnStats()); + if (stats == null) { + return; + } + + CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats serverStatsAggr = null; + try { + serverStatsAggr = + Entity.getCollectionDAO() + .searchIndexServerStatsDAO() + .getAggregatedStats(distributedJob.getId().toString()); + } catch (Exception e) { + LOG.debug("Could not fetch aggregated server stats for job {}", distributedJob.getId(), e); + } + + long successRecords; + long failedRecords; + String statsSource; + + if (serverStatsAggr != null && serverStatsAggr.sinkSuccess() > 0) { + successRecords = serverStatsAggr.sinkSuccess(); + failedRecords = + serverStatsAggr.readerFailed() + + serverStatsAggr.sinkFailed() + + serverStatsAggr.processFailed(); + statsSource = "serverStatsTable"; + } else if (actualSinkStats != null) { + successRecords = actualSinkStats.getSuccessRecords(); + failedRecords = actualSinkStats.getFailedRecords(); + statsSource = "localSink"; + } else { + successRecords = distributedJob.getSuccessRecords(); + failedRecords = distributedJob.getFailedRecords(); + statsSource = "partition-based"; + } + + LOG.debug( + "Stats source: {}, success={}, failed={}", statsSource, successRecords, failedRecords); + + StepStats jobStats = stats.getJobStats(); + if (jobStats != null) { + jobStats.setSuccessRecords(saturatedToInt(successRecords)); + jobStats.setFailedRecords(saturatedToInt(failedRecords)); + } + + StepStats readerStats = stats.getReaderStats(); + if (readerStats != null) { + readerStats.setTotalRecords(saturatedToInt(distributedJob.getTotalRecords())); + long readerFailed = serverStatsAggr != null ? serverStatsAggr.readerFailed() : 0; + long readerWarnings = serverStatsAggr != null ? serverStatsAggr.readerWarnings() : 0; + long readerSuccess = + serverStatsAggr != null + ? serverStatsAggr.readerSuccess() + : distributedJob.getTotalRecords() - readerFailed - readerWarnings; + readerStats.setSuccessRecords(saturatedToInt(readerSuccess)); + readerStats.setFailedRecords(saturatedToInt(readerFailed)); + readerStats.setWarningRecords(saturatedToInt(readerWarnings)); + // Carry stage timing forward into the final ExecutionResult stats. Without this the + // periodic aggregator's totalTimeMs (visible while running) gets clobbered to 0 here, + // and OmAppJobListener picks up the zero on the SUCCESS transition. + if (serverStatsAggr != null) { + readerStats.setTotalTimeMs(serverStatsAggr.readerTimeMs()); + } + } + + StepStats processStats = stats.getProcessStats(); + if (processStats != null && serverStatsAggr != null) { + long processSuccess = serverStatsAggr.processSuccess(); + long processFailed = serverStatsAggr.processFailed(); + processStats.setTotalRecords(saturatedToInt(processSuccess + processFailed)); + processStats.setSuccessRecords(saturatedToInt(processSuccess)); + processStats.setFailedRecords(saturatedToInt(processFailed)); + processStats.setTotalTimeMs(serverStatsAggr.processTimeMs()); + } + + StepStats sinkStats = stats.getSinkStats(); + if (sinkStats != null) { + if (serverStatsAggr != null) { + long sinkSuccess = serverStatsAggr.sinkSuccess(); + long sinkFailed = serverStatsAggr.sinkFailed(); + long actualSinkTotal = sinkSuccess + sinkFailed; + sinkStats.setTotalRecords(saturatedToInt(actualSinkTotal)); + sinkStats.setSuccessRecords(saturatedToInt(sinkSuccess)); + sinkStats.setFailedRecords(saturatedToInt(sinkFailed)); + sinkStats.setTotalTimeMs(serverStatsAggr.sinkTimeMs()); + } else { + long sinkTotal = distributedJob.getTotalRecords(); + sinkStats.setTotalRecords(saturatedToInt(sinkTotal)); + sinkStats.setSuccessRecords(saturatedToInt(successRecords)); + sinkStats.setFailedRecords(saturatedToInt(failedRecords)); + } + } + + StepStats vectorStats = stats.getVectorStats(); + if (vectorStats != null && serverStatsAggr != null) { + long vectorSuccess = serverStatsAggr.vectorSuccess(); + long vectorFailed = serverStatsAggr.vectorFailed(); + vectorStats.setTotalRecords(saturatedToInt(vectorSuccess + vectorFailed)); + vectorStats.setSuccessRecords(saturatedToInt(vectorSuccess)); + vectorStats.setFailedRecords(saturatedToInt(vectorFailed)); + vectorStats.setTotalTimeMs(serverStatsAggr.vectorTimeMs()); + } + + if (distributedJob.getEntityStats() != null && stats.getEntityStats() != null) { + for (Map.Entry entry : + distributedJob.getEntityStats().entrySet()) { + StepStats entityStats = + stats.getEntityStats().getAdditionalProperties().get(entry.getKey()); + if (entityStats != null) { + entityStats.setSuccessRecords(saturatedToInt(entry.getValue().getSuccessRecords())); + entityStats.setFailedRecords(saturatedToInt(entry.getValue().getFailedRecords())); + // Surface all four stage timings on the entity-level StepStats so the UI per-entity + // table can show Reader / Process / Sink / Vector avg latencies side-by-side. + entityStats.setReaderTimeMs(entry.getValue().getReaderTimeMs()); + entityStats.setProcessTimeMs(entry.getValue().getProcessTimeMs()); + entityStats.setSinkTimeMs(entry.getValue().getSinkTimeMs()); + entityStats.setVectorTimeMs(entry.getValue().getVectorTimeMs()); + } + } + } + + updateColumnStatsFromSink(stats); + + StatsReconciler.reconcile(stats); + } + + private void updateColumnStatsFromSink(Stats jobDataStats) { + if (searchIndexSink == null || jobDataStats == null || jobDataStats.getEntityStats() == null) { + return; + } + StepStats columnStats = searchIndexSink.getColumnStats(); + if (columnStats != null) { + StepStats existingColumnStats = + jobDataStats.getEntityStats().getAdditionalProperties().get(Entity.TABLE_COLUMN); + if (existingColumnStats != null) { + existingColumnStats.setTotalRecords(columnStats.getTotalRecords()); + existingColumnStats.setSuccessRecords(columnStats.getSuccessRecords()); + existingColumnStats.setFailedRecords(columnStats.getFailedRecords()); + } + } + } + + private void promoteColumnIndex( + RecreateIndexHandler recreateIndexHandler, + ReindexContext recreateContext, + boolean tableSuccess) { + Optional columnStagedIndex = recreateContext.getStagedIndex(Entity.TABLE_COLUMN); + if (columnStagedIndex.isEmpty()) { + return; + } + try { + finalizeEntityReindex( + recreateIndexHandler, recreateContext, Entity.TABLE_COLUMN, tableSuccess); + LOG.info("Promoted column index (tableSuccess={})", tableSuccess); + } catch (Exception ex) { + LOG.error("Failed to promote column index", ex); + } + } + + private static int saturatedToInt(long value) { + return (int) Math.min(value, Integer.MAX_VALUE); } private ExecutionResult.Status determineStatus(Stats stats) { @@ -295,83 +474,121 @@ private boolean hasIncompleteProcessing(Stats stats) { } private boolean finalizeAllEntityReindex( - RecreateIndexHandler indexPromotionHandler, - ReindexContext stagedIndexContext, + RecreateIndexHandler recreateIndexHandler, + ReindexContext recreateContext, boolean finalSuccess) { - if (indexPromotionHandler == null || stagedIndexContext == null) { + if (recreateIndexHandler == null || recreateContext == null) { return finalSuccess; } - return new DistributedReindexFinalizer(indexPromotionHandler, stagedIndexContext) - .finalizeRemainingEntities(getPromotedEntities(), getFinalEntityStats(), finalSuccess); - } - - private StepStats getColumnStats() { - return searchIndexSink != null ? searchIndexSink.getColumnStats() : null; - } - - private Set getPromotedEntities() { + Set promotedEntities = Collections.emptySet(); if (distributedExecutor != null && distributedExecutor.getEntityTracker() != null) { - return distributedExecutor.getEntityTracker().getPromotedEntities(); + promotedEntities = distributedExecutor.getEntityTracker().getPromotedEntities(); } - return Collections.emptySet(); - } - private Map getFinalEntityStats() { - Map finalEntityStats = new HashMap<>(); - if (distributedExecutor == null) { - mergeInitializedEntityStats(finalEntityStats); - return finalEntityStats; + // Get per-entity stats for determining per-entity success + Map entityStatsMap = Collections.emptyMap(); + if (distributedExecutor != null) { + SearchIndexJob finalJob = distributedExecutor.getJobWithFreshStats(); + if (finalJob != null && finalJob.getEntityStats() != null) { + entityStatsMap = finalJob.getEntityStats(); + } } - SearchIndexJob finalJob = distributedExecutor.getJobWithFreshStats(); - if (finalJob != null && finalJob.getEntityStats() != null) { - finalEntityStats.putAll(finalJob.getEntityStats()); + + LOG.debug( + "Finalization: finalSuccess={}, promotedEntities={}, allEntities={}", + finalSuccess, + promotedEntities, + recreateContext.getEntities()); + + Set entitiesToFinalize = new HashSet<>(recreateContext.getEntities()); + entitiesToFinalize.removeAll(promotedEntities); + + if (promotedEntities.contains(Entity.TABLE) + && !promotedEntities.contains(Entity.TABLE_COLUMN)) { + boolean tableSuccess = computeEntitySuccess(Entity.TABLE, entityStatsMap); + promoteColumnIndex(recreateIndexHandler, recreateContext, tableSuccess); + entitiesToFinalize.remove(Entity.TABLE_COLUMN); } - mergeInitializedEntityStats(finalEntityStats); - return finalEntityStats; - } - private void mergeInitializedEntityStats( - Map finalEntityStats) { - Stats stats = currentStats.get(); - if (stats == null - || stats.getEntityStats() == null - || stats.getEntityStats().getAdditionalProperties() == null) { - return; + LOG.debug("Entities to finalize={}, already promoted={}", entitiesToFinalize, promotedEntities); + + try { + if (!entitiesToFinalize.isEmpty()) { + LOG.info( + "Finalizing {} remaining entities (already promoted: {})", + entitiesToFinalize.size(), + promotedEntities.size()); + + for (String entityType : entitiesToFinalize) { + try { + boolean entitySuccess = computeEntitySuccess(entityType, entityStatsMap); + LOG.debug( + "Finalizing entity '{}' with perEntitySuccess={} (globalSuccess={})", + entityType, + entitySuccess, + finalSuccess); + finalizeEntityReindex(recreateIndexHandler, recreateContext, entityType, entitySuccess); + if (Entity.TABLE.equals(entityType)) { + promoteColumnIndex(recreateIndexHandler, recreateContext, entitySuccess); + } + } catch (Exception ex) { + LOG.error("Failed to finalize reindex for entity: {}", entityType, ex); + } + } + } + } catch (Exception e) { + LOG.error("Error during entity finalization", e); } - stats - .getEntityStats() - .getAdditionalProperties() - .forEach( - (entityType, stepStats) -> - finalEntityStats.computeIfAbsent( - entityType, key -> toEntityTypeStats(key, stepStats))); + return finalSuccess; } - private SearchIndexJob.EntityTypeStats toEntityTypeStats(String entityType, StepStats stepStats) { - long success = stepStats != null ? statValue(stepStats.getSuccessRecords()) : 0L; - long failed = stepStats != null ? statValue(stepStats.getFailedRecords()) : 0L; - return SearchIndexJob.EntityTypeStats.builder() - .entityType(entityType) - .totalRecords(stepStats != null ? statValue(stepStats.getTotalRecords()) : 0L) - .processedRecords(success + failed) - .successRecords(success) - .failedRecords(failed) - .totalPartitions(0) - .completedPartitions(0) - .failedPartitions(0) - .build(); + private boolean computeEntitySuccess( + String entityType, Map entityStatsMap) { + if (entityStatsMap == null || entityStatsMap.isEmpty()) { + return false; + } + SearchIndexJob.EntityTypeStats stats = entityStatsMap.get(entityType); + if (stats == null) { + // Entity not in stats means 0 records — nothing to index = success + return true; + } + return stats.getFailedRecords() == 0 + && stats.getSuccessRecords() + stats.getFailedRecords() >= stats.getTotalRecords(); } - private long statValue(Number value) { - return value != null ? value.longValue() : 0L; + private void finalizeEntityReindex( + RecreateIndexHandler recreateIndexHandler, + ReindexContext recreateContext, + String entityType, + boolean success) { + try { + var entityReindexContext = + org.openmetadata.service.search.EntityReindexContext.builder() + .entityType(entityType) + .originalIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .canonicalIndex(recreateContext.getCanonicalIndex(entityType).orElse(null)) + .activeIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .stagedIndex(recreateContext.getStagedIndex(entityType).orElse(null)) + .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) + .existingAliases(recreateContext.getExistingAliases(entityType)) + .parentAliases( + new HashSet<>(listOrEmpty(recreateContext.getParentAliases(entityType)))) + .build(); + + recreateIndexHandler.finalizeReindex(entityReindexContext, success); + } catch (Exception ex) { + LOG.error("Failed to finalize index recreation flow for {}", entityType, ex); + } } + @Override public Optional getStats() { return Optional.ofNullable(currentStats.get()); } + @Override public void stop() { if (stopped.compareAndSet(false, true)) { LOG.info("Stopping distributed indexing strategy"); @@ -389,6 +606,7 @@ public void stop() { } } + @Override public boolean isStopped() { return stopped.get(); } @@ -449,9 +667,9 @@ Stats initializeTotalRecords(Set entities) { private int getEntityTotal(String entityType) { try { - String correctedType = SearchIndexEntityTypes.normalizeEntityType(entityType); + String correctedType = "queryCostResult".equals(entityType) ? QUERY_COST_RECORD : entityType; - if (!SearchIndexEntityTypes.isTimeSeriesEntity(correctedType)) { + if (!TIME_SERIES_ENTITIES.contains(correctedType)) { return Entity.getEntityRepository(correctedType) .getDao() .listCount(new ListFilter(Include.ALL)); @@ -459,7 +677,7 @@ private int getEntityTotal(String entityType) { ListFilter listFilter = new ListFilter(null); EntityTimeSeriesRepository repository; - if (SearchIndexEntityTypes.isDataInsightEntity(correctedType)) { + if (isDataInsightIndex(correctedType)) { listFilter.addQueryParam("entityFQNHash", FullyQualifiedName.buildHash(correctedType)); repository = Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA); } else { @@ -481,6 +699,10 @@ private int getEntityTotal(String entityType) { } } + private boolean isDataInsightIndex(String entityType) { + return entityType.endsWith("ReportData"); + } + DistributedSearchIndexExecutor getDistributedExecutor() { return distributedExecutor; } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexFinalizer.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexFinalizer.java deleted file mode 100644 index d4aaadd06527..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexFinalizer.java +++ /dev/null @@ -1,143 +0,0 @@ -/* - * Copyright 2024 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.service.apps.bundles.searchIndex; - -import java.util.HashSet; -import java.util.Map; -import java.util.Set; -import lombok.extern.slf4j.Slf4j; -import org.openmetadata.service.Entity; -import org.openmetadata.service.apps.bundles.searchIndex.distributed.SearchIndexJob; -import org.openmetadata.service.search.RecreateIndexHandler; -import org.openmetadata.service.search.ReindexContext; - -@Slf4j -class DistributedReindexFinalizer { - private final RecreateIndexHandler indexPromotionHandler; - private final ReindexContext stagedIndexContext; - - DistributedReindexFinalizer( - RecreateIndexHandler indexPromotionHandler, ReindexContext stagedIndexContext) { - this.indexPromotionHandler = indexPromotionHandler; - this.stagedIndexContext = stagedIndexContext; - } - - boolean finalizeRemainingEntities( - Set promotedEntities, - Map entityStats, - boolean finalSuccess) { - LOG.debug( - "Finalization: finalSuccess={}, promotedEntities={}, allEntities={}", - finalSuccess, - promotedEntities, - stagedIndexContext.getEntities()); - - Set entitiesToFinalize = new HashSet<>(stagedIndexContext.getEntities()); - entitiesToFinalize.removeAll(promotedEntities); - Set finalizedEntities = new HashSet<>(promotedEntities); - - routeColumnFinalizationThroughTable(entitiesToFinalize); - promoteColumnIndexIfTableWasPromoted( - promotedEntities, entityStats, entitiesToFinalize, finalizedEntities); - finalizeEntities(entitiesToFinalize, entityStats, finalSuccess, finalizedEntities); - - return finalSuccess; - } - - private void routeColumnFinalizationThroughTable(Set entitiesToFinalize) { - if (entitiesToFinalize.contains(Entity.TABLE)) { - entitiesToFinalize.remove(Entity.TABLE_COLUMN); - } - } - - private void promoteColumnIndexIfTableWasPromoted( - Set promotedEntities, - Map entityStats, - Set entitiesToFinalize, - Set finalizedEntities) { - if (promotedEntities.contains(Entity.TABLE) - && !promotedEntities.contains(Entity.TABLE_COLUMN)) { - boolean tableSuccess = computeEntitySuccess(Entity.TABLE, entityStats); - promoteColumnIndex(tableSuccess, finalizedEntities); - entitiesToFinalize.remove(Entity.TABLE_COLUMN); - } - } - - private void finalizeEntities( - Set entitiesToFinalize, - Map entityStats, - boolean finalSuccess, - Set finalizedEntities) { - LOG.debug("Entities to finalize={}", entitiesToFinalize); - if (entitiesToFinalize.isEmpty()) { - return; - } - - LOG.info("Finalizing {} remaining entities", entitiesToFinalize.size()); - for (String entityType : entitiesToFinalize) { - if (!finalizedEntities.add(entityType)) { - LOG.debug("Skipping already finalized entity '{}'", entityType); - continue; - } - try { - boolean entitySuccess = computeEntitySuccess(entityType, entityStats); - LOG.debug( - "Finalizing entity '{}' with perEntitySuccess={} (globalSuccess={})", - entityType, - entitySuccess, - finalSuccess); - finalizeEntityReindex(entityType, entitySuccess); - if (Entity.TABLE.equals(entityType)) { - promoteColumnIndex(entitySuccess, finalizedEntities); - } - } catch (Exception ex) { - LOG.error("Failed to finalize reindex for entity: {}", entityType, ex); - } - } - } - - private void promoteColumnIndex(boolean tableSuccess, Set finalizedEntities) { - if (stagedIndexContext.getStagedIndex(Entity.TABLE_COLUMN).isEmpty()) { - return; - } - if (!finalizedEntities.add(Entity.TABLE_COLUMN)) { - LOG.debug("Skipping already finalized column index"); - return; - } - try { - finalizeEntityReindex(Entity.TABLE_COLUMN, tableSuccess); - LOG.info("Promoted column index (tableSuccess={})", tableSuccess); - } catch (Exception ex) { - LOG.error("Failed to promote column index", ex); - } - } - - private boolean computeEntitySuccess( - String entityType, Map entityStats) { - if (entityStats == null || entityStats.isEmpty()) { - return false; - } - SearchIndexJob.EntityTypeStats stats = entityStats.get(entityType); - if (stats == null) { - return false; - } - return stats.getFailedRecords() == 0 - && stats.getSuccessRecords() + stats.getFailedRecords() >= stats.getTotalRecords(); - } - - private void finalizeEntityReindex(String entityType, boolean success) { - indexPromotionHandler.finalizeReindex( - EntityReindexContextMapper.fromStagedContext(stagedIndexContext, entityType), success); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexStatsMapper.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexStatsMapper.java deleted file mode 100644 index 6d29711d08af..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexStatsMapper.java +++ /dev/null @@ -1,234 +0,0 @@ -/* - * Copyright 2024 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.service.apps.bundles.searchIndex; - -import java.util.Map; -import lombok.extern.slf4j.Slf4j; -import org.openmetadata.schema.system.Stats; -import org.openmetadata.schema.system.StepStats; -import org.openmetadata.service.Entity; -import org.openmetadata.service.apps.bundles.searchIndex.distributed.SearchIndexJob; -import org.openmetadata.service.jdbi3.CollectionDAO; - -@Slf4j -class DistributedReindexStatsMapper { - private final CollectionDAO collectionDAO; - - DistributedReindexStatsMapper(CollectionDAO collectionDAO) { - this.collectionDAO = collectionDAO; - } - - void updateStats( - Stats stats, - SearchIndexJob distributedJob, - StepStats actualSinkStats, - StepStats columnStats) { - if (stats == null) { - return; - } - - CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats aggregatedStats = - getAggregatedServerStats(distributedJob); - StatsSource source = resolveStatsSource(distributedJob, aggregatedStats, actualSinkStats); - - LOG.debug( - "Stats source: {}, success={}, failed={}", - source.name(), - source.successRecords(), - source.failedRecords()); - - updateJobStats(stats, source); - updateReaderStats(stats, distributedJob, aggregatedStats); - updateProcessStats(stats, aggregatedStats); - updateSinkStats(stats, distributedJob, aggregatedStats, source); - updateVectorStats(stats, aggregatedStats); - updateEntityStats(stats, distributedJob); - updateColumnStats(stats, columnStats); - - StatsReconciler.reconcile(stats); - } - - private CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats getAggregatedServerStats( - SearchIndexJob distributedJob) { - try { - return collectionDAO - .searchIndexServerStatsDAO() - .getAggregatedStats(distributedJob.getId().toString()); - } catch (Exception e) { - LOG.debug("Could not fetch aggregated server stats for job {}", distributedJob.getId(), e); - return null; - } - } - - private StatsSource resolveStatsSource( - SearchIndexJob distributedJob, - CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats aggregatedStats, - StepStats actualSinkStats) { - if (hasAggregatedStageRecords(aggregatedStats)) { - return new StatsSource( - "serverStatsTable", - aggregatedStats.sinkSuccess(), - aggregatedStats.readerFailed() - + aggregatedStats.sinkFailed() - + aggregatedStats.processFailed()); - } - if (actualSinkStats != null) { - return new StatsSource( - "localSink", actualSinkStats.getSuccessRecords(), actualSinkStats.getFailedRecords()); - } - return new StatsSource( - "partition-based", distributedJob.getSuccessRecords(), distributedJob.getFailedRecords()); - } - - private boolean hasAggregatedStageRecords( - CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats aggregatedStats) { - return aggregatedStats != null - && (aggregatedStats.readerSuccess() > 0 - || aggregatedStats.readerFailed() > 0 - || aggregatedStats.readerWarnings() > 0 - || aggregatedStats.processSuccess() > 0 - || aggregatedStats.processFailed() > 0 - || aggregatedStats.sinkSuccess() > 0 - || aggregatedStats.sinkFailed() > 0 - || aggregatedStats.vectorSuccess() > 0 - || aggregatedStats.vectorFailed() > 0); - } - - private void updateJobStats(Stats stats, StatsSource source) { - StepStats jobStats = stats.getJobStats(); - if (jobStats != null) { - jobStats.setSuccessRecords(saturatedToInt(source.successRecords())); - jobStats.setFailedRecords(saturatedToInt(source.failedRecords())); - } - } - - private void updateReaderStats( - Stats stats, - SearchIndexJob distributedJob, - CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats aggregatedStats) { - StepStats readerStats = stats.getReaderStats(); - if (readerStats == null) { - return; - } - - readerStats.setTotalRecords(saturatedToInt(distributedJob.getTotalRecords())); - long readerFailed = aggregatedStats != null ? aggregatedStats.readerFailed() : 0; - long readerWarnings = aggregatedStats != null ? aggregatedStats.readerWarnings() : 0; - long readerSuccess = - aggregatedStats != null - ? aggregatedStats.readerSuccess() - : distributedJob.getTotalRecords() - readerFailed - readerWarnings; - readerStats.setSuccessRecords(saturatedToInt(readerSuccess)); - readerStats.setFailedRecords(saturatedToInt(readerFailed)); - readerStats.setWarningRecords(saturatedToInt(readerWarnings)); - if (aggregatedStats != null) { - readerStats.setTotalTimeMs(aggregatedStats.readerTimeMs()); - } - } - - private void updateProcessStats( - Stats stats, CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats aggregatedStats) { - StepStats processStats = stats.getProcessStats(); - if (processStats == null || aggregatedStats == null) { - return; - } - - long processSuccess = aggregatedStats.processSuccess(); - long processFailed = aggregatedStats.processFailed(); - processStats.setTotalRecords(saturatedToInt(processSuccess + processFailed)); - processStats.setSuccessRecords(saturatedToInt(processSuccess)); - processStats.setFailedRecords(saturatedToInt(processFailed)); - processStats.setTotalTimeMs(aggregatedStats.processTimeMs()); - } - - private void updateSinkStats( - Stats stats, - SearchIndexJob distributedJob, - CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats aggregatedStats, - StatsSource source) { - StepStats sinkStats = stats.getSinkStats(); - if (sinkStats == null) { - return; - } - - if (aggregatedStats != null) { - long sinkSuccess = aggregatedStats.sinkSuccess(); - long sinkFailed = aggregatedStats.sinkFailed(); - sinkStats.setTotalRecords(saturatedToInt(sinkSuccess + sinkFailed)); - sinkStats.setSuccessRecords(saturatedToInt(sinkSuccess)); - sinkStats.setFailedRecords(saturatedToInt(sinkFailed)); - sinkStats.setTotalTimeMs(aggregatedStats.sinkTimeMs()); - return; - } - - sinkStats.setTotalRecords(saturatedToInt(distributedJob.getTotalRecords())); - sinkStats.setSuccessRecords(saturatedToInt(source.successRecords())); - sinkStats.setFailedRecords(saturatedToInt(source.failedRecords())); - } - - private void updateVectorStats( - Stats stats, CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats aggregatedStats) { - StepStats vectorStats = stats.getVectorStats(); - if (vectorStats == null || aggregatedStats == null) { - return; - } - - long vectorSuccess = aggregatedStats.vectorSuccess(); - long vectorFailed = aggregatedStats.vectorFailed(); - vectorStats.setTotalRecords(saturatedToInt(vectorSuccess + vectorFailed)); - vectorStats.setSuccessRecords(saturatedToInt(vectorSuccess)); - vectorStats.setFailedRecords(saturatedToInt(vectorFailed)); - vectorStats.setTotalTimeMs(aggregatedStats.vectorTimeMs()); - } - - private void updateEntityStats(Stats stats, SearchIndexJob distributedJob) { - if (distributedJob.getEntityStats() == null || stats.getEntityStats() == null) { - return; - } - - for (Map.Entry entry : - distributedJob.getEntityStats().entrySet()) { - StepStats entityStats = stats.getEntityStats().getAdditionalProperties().get(entry.getKey()); - if (entityStats != null) { - SearchIndexJob.EntityTypeStats distributedEntityStats = entry.getValue(); - entityStats.setSuccessRecords(saturatedToInt(distributedEntityStats.getSuccessRecords())); - entityStats.setFailedRecords(saturatedToInt(distributedEntityStats.getFailedRecords())); - entityStats.setReaderTimeMs(distributedEntityStats.getReaderTimeMs()); - entityStats.setProcessTimeMs(distributedEntityStats.getProcessTimeMs()); - entityStats.setSinkTimeMs(distributedEntityStats.getSinkTimeMs()); - entityStats.setVectorTimeMs(distributedEntityStats.getVectorTimeMs()); - } - } - } - - private void updateColumnStats(Stats stats, StepStats columnStats) { - if (columnStats == null || stats.getEntityStats() == null) { - return; - } - - StepStats existingColumnStats = - stats.getEntityStats().getAdditionalProperties().get(Entity.TABLE_COLUMN); - if (existingColumnStats != null) { - existingColumnStats.setTotalRecords(columnStats.getTotalRecords()); - existingColumnStats.setSuccessRecords(columnStats.getSuccessRecords()); - existingColumnStats.setFailedRecords(columnStats.getFailedRecords()); - } - } - - private static int saturatedToInt(long value) { - return (int) Math.min(value, Integer.MAX_VALUE); - } - - private record StatsSource(String name, long successRecords, long failedRecords) {} -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimator.java new file mode 100644 index 000000000000..512d990ee2a7 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimator.java @@ -0,0 +1,38 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Set; + +/** + * Per-entity-type batch sizing based on typical document size. Large entity types (tables, + * dashboards, etc.) produce bigger search documents, so we use smaller batches. Small entity types + * (users, tags, etc.) produce tiny documents, so we can use larger batches. + */ +public final class EntityBatchSizeEstimator { + + private static final Set LARGE_ENTITIES = + Set.of("table", "topic", "dashboard", "mlmodel", "container", "storedProcedure"); + + private static final Set SMALL_ENTITIES = + Set.of("user", "team", "bot", "role", "policy", "tag", "classification"); + + private static final int MIN_BATCH_SIZE = 25; + private static final int MAX_BATCH_SIZE = 1000; + + private EntityBatchSizeEstimator() {} + + public static int estimateBatchSize(String entityType, int baseBatchSize) { + if (baseBatchSize <= 0) { + return baseBatchSize; + } + + if (LARGE_ENTITIES.contains(entityType)) { + return Math.max(baseBatchSize / 2, MIN_BATCH_SIZE); + } + + if (SMALL_ENTITIES.contains(entityType)) { + return Math.min(baseBatchSize * 2, MAX_BATCH_SIZE); + } + + return baseBatchSize; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReader.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReader.java new file mode 100644 index 000000000000..6ae3400d235f --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReader.java @@ -0,0 +1,354 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.service.Entity.QUERY_COST_RECORD; +import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; +import static org.openmetadata.service.Entity.TEST_CASE_RESULT; +import static org.openmetadata.service.workflows.searchIndex.ReindexingUtil.getSearchIndexFields; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Phaser; +import java.util.concurrent.atomic.AtomicBoolean; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.exception.SearchIndexException; +import org.openmetadata.service.util.RestUtil; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntitiesSource; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntityTimeSeriesSource; + +/** + * Standalone reader that encapsulates all entity reading logic. Decoupled from queues and sinks — + * delivers batches via a callback interface. + */ +@Slf4j +public class EntityReader implements AutoCloseable { + + static final Set TIME_SERIES_ENTITIES = + Set.of( + ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), + ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), + ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), + TEST_CASE_RESOLUTION_STATUS, + TEST_CASE_RESULT, + QUERY_COST_RECORD); + + private static final int MAX_READERS_PER_ENTITY = 5; + + @FunctionalInterface + public interface BatchCallback { + void onBatchRead(String entityType, ResultList batch, int offset) + throws InterruptedException; + } + + @FunctionalInterface + interface KeysetBatchReader { + ResultList readNextKeyset(String cursor) throws SearchIndexException; + } + + @FunctionalInterface + interface BoundaryFinder { + List findBoundaries(int numReaders, int totalRecords); + } + + private static final int DEFAULT_MAX_RETRY_ATTEMPTS = 3; + private static final long DEFAULT_RETRY_BACKOFF_MS = 500; + + private final ExecutorService producerExecutor; + private final AtomicBoolean stopped; + private final int maxRetryAttempts; + private final long retryBackoffMs; + + public EntityReader(ExecutorService producerExecutor, AtomicBoolean stopped) { + this(producerExecutor, stopped, DEFAULT_MAX_RETRY_ATTEMPTS, DEFAULT_RETRY_BACKOFF_MS); + } + + public EntityReader( + ExecutorService producerExecutor, + AtomicBoolean stopped, + int maxRetryAttempts, + long retryBackoffMs) { + this.producerExecutor = producerExecutor; + this.stopped = stopped; + this.maxRetryAttempts = maxRetryAttempts; + this.retryBackoffMs = retryBackoffMs; + } + + /** + * Read all entities of a given type, invoking callback for each batch. + * + * @param entityType The entity type to read + * @param totalRecords Total records expected for this entity + * @param batchSize Batch size for reading + * @param phaser Phaser for completion tracking (readers will register/deregister) + * @param callback Callback invoked with each batch + * @return Number of readers submitted + */ + public int readEntity( + String entityType, int totalRecords, int batchSize, Phaser phaser, BatchCallback callback) { + return readEntity(entityType, totalRecords, batchSize, phaser, callback, null, null); + } + + public int readEntity( + String entityType, + int totalRecords, + int batchSize, + Phaser phaser, + BatchCallback callback, + Long timeSeriesStartTs, + Long timeSeriesEndTs) { + if (totalRecords <= 0) { + return 0; + } + + int numReaders = + Math.min(calculateNumberOfReaders(totalRecords, batchSize), MAX_READERS_PER_ENTITY); + phaser.bulkRegister(numReaders); + + try { + if (TIME_SERIES_ENTITIES.contains(entityType)) { + submitReaders( + entityType, + totalRecords, + batchSize, + numReaders, + phaser, + callback, + () -> { + PaginatedEntityTimeSeriesSource source = + (timeSeriesStartTs != null) + ? new PaginatedEntityTimeSeriesSource( + entityType, + batchSize, + getSearchIndexFields(entityType), + totalRecords, + timeSeriesStartTs, + timeSeriesEndTs) + : new PaginatedEntityTimeSeriesSource( + entityType, batchSize, getSearchIndexFields(entityType), totalRecords); + return source::readWithCursor; + }, + (readers, total) -> { + List cursors = new ArrayList<>(); + int perReader = total / readers; + for (int i = 1; i < readers; i++) { + cursors.add(RestUtil.encodeCursor(String.valueOf(i * perReader))); + } + return cursors; + }); + } else { + PaginatedEntitiesSource entSource = + new PaginatedEntitiesSource( + entityType, batchSize, getSearchIndexFields(entityType), totalRecords); + submitReaders( + entityType, + totalRecords, + batchSize, + numReaders, + phaser, + callback, + () -> { + PaginatedEntitiesSource source = + new PaginatedEntitiesSource( + entityType, batchSize, getSearchIndexFields(entityType), totalRecords); + return source::readNextKeyset; + }, + entSource::findBoundaryCursors); + } + } catch (Exception e) { + LOG.error( + "Failed to submit readers for {}, deregistering {} phaser parties", + entityType, + numReaders, + e); + for (int i = 0; i < numReaders; i++) { + phaser.arriveAndDeregister(); + } + throw e; + } + + return numReaders; + } + + public void stop() { + stopped.set(true); + } + + @Override + public void close() { + stop(); + } + + private void submitReaders( + String entityType, + int totalRecords, + int batchSize, + int numReaders, + Phaser phaser, + BatchCallback callback, + java.util.function.Supplier readerFactory, + BoundaryFinder boundaryFinder) { + if (numReaders == 1) { + KeysetBatchReader reader = readerFactory.get(); + producerExecutor.submit( + () -> + readKeysetBatches( + entityType, Integer.MAX_VALUE, batchSize, null, reader, phaser, callback)); + return; + } + + List boundaries = boundaryFinder.findBoundaries(numReaders, totalRecords); + int actualReaders = boundaries.size() + 1; + int recordsPerReader = (totalRecords + actualReaders - 1) / actualReaders; + + if (actualReaders < numReaders) { + LOG.warn( + "Boundary discovery for {} returned {} cursors (expected {}), using {} readers", + entityType, + boundaries.size(), + numReaders - 1, + actualReaders); + for (int j = 0; j < numReaders - actualReaders; j++) { + phaser.arriveAndDeregister(); + } + } + + for (int i = 0; i < actualReaders; i++) { + String startCursor = (i == 0) ? null : boundaries.get(i - 1); + int limit = (i == actualReaders - 1) ? Integer.MAX_VALUE : recordsPerReader; + KeysetBatchReader readerSource = readerFactory.get(); + final int readerLimit = limit; + producerExecutor.submit( + () -> + readKeysetBatches( + entityType, readerLimit, batchSize, startCursor, readerSource, phaser, callback)); + } + } + + private void readKeysetBatches( + String entityType, + int recordLimit, + int batchSize, + String startCursor, + KeysetBatchReader batchReader, + Phaser phaser, + BatchCallback callback) { + // Bypass the Redis-backed entity cache on the reader thread for the same reasons as the + // distributed PartitionWorker: bulk reindex never re-reads entities, every relationship + // lookup pays a cache round-trip we don't need, and an unhealthy Redis turns each lookup + // into a 300ms timeout. See {@link org.openmetadata.service.cache.EntityCacheBypass}. + try (org.openmetadata.service.cache.EntityCacheBypass.Handle ignored = + org.openmetadata.service.cache.EntityCacheBypass.skip()) { + readKeysetBatchesInternal( + entityType, recordLimit, batchSize, startCursor, batchReader, phaser, callback); + } + } + + private void readKeysetBatchesInternal( + String entityType, + int recordLimit, + int batchSize, + String startCursor, + KeysetBatchReader batchReader, + Phaser phaser, + BatchCallback callback) { + try { + String keysetCursor = startCursor; + int processed = 0; + + while (processed < recordLimit && !stopped.get()) { + ResultList result = readWithRetry(batchReader, keysetCursor, entityType); + if (stopped.get()) { + break; + } + + if (result == null || result.getData().isEmpty()) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (empty result)", + entityType, + processed, + recordLimit); + break; + } + + callback.onBatchRead(entityType, result, processed); + + int readCount = result.getData().size(); + int errorCount = result.getErrors() != null ? result.getErrors().size() : 0; + int warningsCount = result.getWarningsCount() != null ? result.getWarningsCount() : 0; + processed += readCount + errorCount + warningsCount; + + keysetCursor = result.getPaging() != null ? result.getPaging().getAfter() : null; + if (keysetCursor == null) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (null cursor)", + entityType, + processed, + recordLimit); + break; + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warn("Interrupted during reading of {}", entityType); + } catch (SearchIndexException e) { + LOG.error("Error reading keyset batch for {}", entityType, e); + } catch (Exception e) { + if (!stopped.get()) { + LOG.error("Error in keyset reading for {}", entityType, e); + } + } finally { + phaser.arriveAndDeregister(); + } + } + + private ResultList readWithRetry( + KeysetBatchReader batchReader, String keysetCursor, String entityType) + throws SearchIndexException, InterruptedException { + for (int attempt = 0; attempt <= maxRetryAttempts; attempt++) { + try { + return batchReader.readNextKeyset(keysetCursor); + } catch (SearchIndexException e) { + if (attempt >= maxRetryAttempts || !isTransientError(e)) { + throw e; + } + long backoff = retryBackoffMs * (1L << attempt); + LOG.warn( + "Transient read failure for {} (attempt {}/{}), retrying in {}ms", + entityType, + attempt + 1, + maxRetryAttempts, + backoff); + Thread.sleep(Math.min(backoff, 10_000)); + } + } + return null; + } + + static boolean isTransientError(SearchIndexException e) { + String msg = e.getMessage(); + if (msg == null) { + return false; + } + String lower = msg.toLowerCase(); + return lower.contains("timeout") + || lower.contains("connection") + || lower.contains("pool exhausted") + || lower.contains("connectexception") + || lower.contains("sockettimeoutexception"); + } + + static List getSearchIndexFields(String entityType) { + return org.openmetadata.service.workflows.searchIndex.ReindexingUtil.getSearchIndexFields( + entityType); + } + + static int calculateNumberOfReaders(int totalEntityRecords, int batchSize) { + if (batchSize <= 0) return 1; + return (totalEntityRecords + batchSize - 1) / batchSize; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReindexContextMapper.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReindexContextMapper.java deleted file mode 100644 index a501b73e27c9..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReindexContextMapper.java +++ /dev/null @@ -1,40 +0,0 @@ -/* - * Copyright 2024 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.service.apps.bundles.searchIndex; - -import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; - -import java.util.HashSet; -import org.openmetadata.service.search.EntityReindexContext; -import org.openmetadata.service.search.ReindexContext; - -public final class EntityReindexContextMapper { - private EntityReindexContextMapper() {} - - public static EntityReindexContext fromStagedContext( - ReindexContext stagedIndexContext, String entityType) { - String originalIndex = stagedIndexContext.getOriginalIndex(entityType).orElse(null); - - return EntityReindexContext.builder() - .entityType(entityType) - .originalIndex(originalIndex) - .canonicalIndex(stagedIndexContext.getCanonicalIndex(entityType).orElse(null)) - .activeIndex(originalIndex) - .stagedIndex(stagedIndexContext.getStagedIndex(entityType).orElse(null)) - .canonicalAliases(stagedIndexContext.getCanonicalAlias(entityType).orElse(null)) - .existingAliases(stagedIndexContext.getExistingAliases(entityType)) - .parentAliases(new HashSet<>(listOrEmpty(stagedIndexContext.getParentAliases(entityType)))) - .build(); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipeline.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipeline.java new file mode 100644 index 000000000000..d45864f497f1 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipeline.java @@ -0,0 +1,603 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; +import static org.openmetadata.service.workflows.searchIndex.ReindexingUtil.isDataInsightIndex; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Phaser; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.Entity; +import org.openmetadata.service.jdbi3.EntityRepository; +import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; +import org.openmetadata.service.jdbi3.ListFilter; +import org.openmetadata.service.search.EntityReindexContext; +import org.openmetadata.service.search.RecreateIndexHandler; +import org.openmetadata.service.search.ReindexContext; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.util.FullyQualifiedName; +import org.openmetadata.service.workflows.searchIndex.ReindexingUtil; +import org.slf4j.MDC; + +/** + * Quartz-decoupled indexing pipeline that orchestrates: entity discovery -> reader -> queue -> sink. + * This class can be used by SearchIndexExecutor, CLI tools, REST APIs, or unit tests. + */ +@Slf4j +public class IndexingPipeline implements AutoCloseable { + + private static final String POISON_PILL = "__POISON_PILL__"; + private static final int DEFAULT_QUEUE_SIZE = 20000; + private static final int MAX_CONSUMER_THREADS = + Math.min(20, Runtime.getRuntime().availableProcessors() * 2); + private static final int MAX_JOB_THREADS = + Math.min(30, Runtime.getRuntime().availableProcessors() * 4); + private static final String ENTITY_TYPE_KEY = "entityType"; + private static final String RECREATE_INDEX = "recreateIndex"; + + private final SearchRepository searchRepository; + private final CompositeProgressListener listeners; + private final AtomicBoolean stopped = new AtomicBoolean(false); + @Getter private final AtomicReference stats = new AtomicReference<>(); + + private BulkSink searchIndexSink; + private RecreateIndexHandler recreateIndexHandler; + private ReindexContext recreateContext; + private EntityReader entityReader; + private ExecutorService consumerExecutor; + private ExecutorService producerExecutor; + private ExecutorService jobExecutor; + private BlockingQueue> taskQueue; + private final Set promotedEntities = java.util.concurrent.ConcurrentHashMap.newKeySet(); + + record IndexingTask(String entityType, ResultList entities, int offset) {} + + public IndexingPipeline(SearchRepository searchRepository) { + this.searchRepository = searchRepository; + this.listeners = new CompositeProgressListener(); + } + + public IndexingPipeline addListener(ReindexingProgressListener listener) { + listeners.addListener(listener); + return this; + } + + public ExecutionResult execute( + ReindexingConfiguration config, + ReindexingJobContext context, + Set entities, + BulkSink sink, + RecreateIndexHandler handler, + ReindexContext recreateCtx) { + this.searchIndexSink = sink; + this.recreateIndexHandler = handler; + this.recreateContext = recreateCtx; + long startTime = System.currentTimeMillis(); + + stats.set(initializeStats(config, entities)); + listeners.onJobStarted(context); + + try { + runPipeline(config, entities); + closeSink(); + finalizeReindex(); + return buildResult(startTime); + } catch (Exception e) { + LOG.error("Pipeline execution failed", e); + listeners.onJobFailed(stats.get(), e); + return ExecutionResult.fromStats(stats.get(), ExecutionResult.Status.FAILED, startTime); + } + } + + private void runPipeline(ReindexingConfiguration config, Set entities) + throws InterruptedException { + int numConsumers = + config.consumerThreads() > 0 ? Math.min(config.consumerThreads(), MAX_CONSUMER_THREADS) : 2; + int queueSize = config.queueSize() > 0 ? config.queueSize() : DEFAULT_QUEUE_SIZE; + int batchSize = config.batchSize(); + + taskQueue = new LinkedBlockingQueue<>(queueSize); + String jobIdTag = MDC.get("reindexJobId"); + String threadPrefix = "reindex-" + (jobIdTag != null ? jobIdTag + "-" : ""); + consumerExecutor = + Executors.newFixedThreadPool( + numConsumers, + Thread.ofPlatform().name(threadPrefix + "pipeline-consumer-", 0).factory()); + producerExecutor = + Executors.newFixedThreadPool( + config.producerThreads() > 0 ? config.producerThreads() : 2, + Thread.ofPlatform().name(threadPrefix + "pipeline-producer-", 0).factory()); + jobExecutor = + Executors.newFixedThreadPool( + Math.min(entities.size(), MAX_JOB_THREADS), + Thread.ofPlatform().name(threadPrefix + "pipeline-job-", 0).factory()); + + entityReader = new EntityReader(producerExecutor, stopped); + + CountDownLatch consumerLatch = new CountDownLatch(numConsumers); + Map mdc = MDC.getCopyOfContextMap(); + for (int i = 0; i < numConsumers; i++) { + final int id = i; + consumerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + runConsumer(id, consumerLatch); + } finally { + MDC.clear(); + } + }); + } + + try { + readAllEntities(config, entities, batchSize); + signalConsumersToStop(numConsumers); + consumerLatch.await(); + } catch (InterruptedException e) { + stopped.set(true); + Thread.currentThread().interrupt(); + throw e; + } finally { + shutdownExecutors(); + } + } + + private void readAllEntities(ReindexingConfiguration config, Set entities, int batchSize) + throws InterruptedException { + List ordered = EntityPriority.sortByPriority(entities); + Phaser producerPhaser = new Phaser(entities.size()); + Map mdc = MDC.getCopyOfContextMap(); + + for (String entityType : ordered) { + jobExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + int totalRecords = getTotalEntityRecords(entityType); + listeners.onEntityTypeStarted(entityType, totalRecords); + + int effectiveBatchSize = + EntityBatchSizeEstimator.estimateBatchSize(entityType, batchSize); + Long filterStartTs = null; + Long filterEndTs = null; + long startTs = config.getTimeSeriesStartTs(entityType); + if (startTs > 0) { + filterStartTs = startTs; + filterEndTs = System.currentTimeMillis(); + } + entityReader.readEntity( + entityType, + totalRecords, + effectiveBatchSize, + producerPhaser, + (type, batch, offset) -> { + if (!stopped.get()) { + taskQueue.put(new IndexingTask<>(type, batch, offset)); + } + }, + filterStartTs, + filterEndTs); + } catch (Exception e) { + LOG.error("Error reading entity type {}", entityType, e); + } finally { + producerPhaser.arriveAndDeregister(); + MDC.clear(); + } + }); + } + + int phase = 0; + while (!producerPhaser.isTerminated()) { + if (stopped.get() || Thread.currentThread().isInterrupted()) { + break; + } + try { + producerPhaser.awaitAdvanceInterruptibly(phase, 1, TimeUnit.SECONDS); + break; + } catch (TimeoutException e) { + // Continue + } + } + } + + @SuppressWarnings("unchecked") + private void runConsumer(int consumerId, CountDownLatch consumerLatch) { + try { + while (!stopped.get()) { + IndexingTask task = taskQueue.poll(200, TimeUnit.MILLISECONDS); + if (task == null) continue; + if (POISON_PILL.equals(task.entityType())) break; + + String entityType = task.entityType(); + ResultList entities = task.entities(); + Map contextData = createContextData(entityType); + + int readerSuccess = listOrEmpty(entities.getData()).size(); + int readerFailed = listOrEmpty(entities.getErrors()).size(); + int readerWarnings = entities.getWarningsCount() != null ? entities.getWarningsCount() : 0; + updateReaderStats(readerSuccess, readerFailed, readerWarnings); + + try { + if (!EntityReader.TIME_SERIES_ENTITIES.contains(entityType)) { + searchIndexSink.write(entities.getData(), contextData); + } else { + searchIndexSink.write(entities.getData(), contextData); + } + + StepStats entityStats = new StepStats(); + entityStats.setSuccessRecords(readerSuccess); + entityStats.setFailedRecords(readerFailed); + updateEntityAndJobStats(entityType, entityStats); + + if (Entity.TABLE.equals(entityType)) { + updateColumnStatsFromSink(); + } + + listeners.onProgressUpdate(stats.get(), null); + } catch (Exception e) { + LOG.error("Sink error for {}", entityType, e); + IndexingError error = + new IndexingError() + .withErrorSource(IndexingError.ErrorSource.SINK) + .withMessage(e.getMessage()); + listeners.onError(entityType, error, stats.get()); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + consumerLatch.countDown(); + } + } + + private Map createContextData(String entityType) { + Map contextData = new HashMap<>(); + contextData.put(ENTITY_TYPE_KEY, entityType); + contextData.put(RECREATE_INDEX, recreateContext != null); + if (recreateContext != null) { + contextData.put(ReindexingUtil.RECREATE_CONTEXT, recreateContext); + recreateContext + .getStagedIndex(entityType) + .ifPresent(index -> contextData.put(ReindexingUtil.TARGET_INDEX_KEY, index)); + } + return contextData; + } + + private void signalConsumersToStop(int numConsumers) throws InterruptedException { + for (int i = 0; i < numConsumers; i++) { + taskQueue.put(new IndexingTask<>(POISON_PILL, null, -1)); + } + } + + private void closeSink() { + if (searchIndexSink != null) { + int pendingVectorTasks = searchIndexSink.getPendingVectorTaskCount(); + if (pendingVectorTasks > 0) { + LOG.info("Waiting for {} pending vector embedding tasks", pendingVectorTasks); + VectorCompletionResult vcResult = searchIndexSink.awaitVectorCompletionWithDetails(300); + LOG.info( + "Vector completion: completed={}, pending={}, waited={}ms", + vcResult.completed(), + vcResult.pendingTaskCount(), + vcResult.waitedMillis()); + } + searchIndexSink.close(); + syncSinkStats(); + } + } + + private void finalizeReindex() { + if (recreateIndexHandler == null || recreateContext == null) return; + + try { + recreateContext + .getEntities() + .forEach( + entityType -> { + if (promotedEntities.contains(entityType)) return; + try { + EntityReindexContext ctx = buildEntityReindexContext(entityType); + recreateIndexHandler.finalizeReindex(ctx, !stopped.get()); + } catch (Exception ex) { + LOG.error("Failed to finalize reindex for {}", entityType, ex); + } + }); + } finally { + recreateContext = null; + promotedEntities.clear(); + } + } + + private EntityReindexContext buildEntityReindexContext(String entityType) { + return EntityReindexContext.builder() + .entityType(entityType) + .originalIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .canonicalIndex(recreateContext.getCanonicalIndex(entityType).orElse(null)) + .activeIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .stagedIndex(recreateContext.getStagedIndex(entityType).orElse(null)) + .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) + .existingAliases(recreateContext.getExistingAliases(entityType)) + .parentAliases( + new HashSet<>( + org.openmetadata.common.utils.CommonUtil.listOrEmpty( + recreateContext.getParentAliases(entityType)))) + .build(); + } + + private ExecutionResult buildResult(long startTime) { + syncSinkStats(); + updateColumnStatsFromSink(); + Stats currentStats = stats.get(); + if (currentStats != null) { + StatsReconciler.reconcile(currentStats); + } + + ExecutionResult.Status status; + if (stopped.get()) { + status = ExecutionResult.Status.STOPPED; + listeners.onJobStopped(currentStats); + } else if (hasFailures()) { + status = ExecutionResult.Status.COMPLETED_WITH_ERRORS; + listeners.onJobCompletedWithErrors(currentStats, System.currentTimeMillis() - startTime); + } else { + status = ExecutionResult.Status.COMPLETED; + listeners.onJobCompleted(currentStats, System.currentTimeMillis() - startTime); + } + + return ExecutionResult.fromStats(currentStats, status, startTime); + } + + private boolean hasFailures() { + Stats s = stats.get(); + if (s == null || s.getJobStats() == null) return false; + StepStats js = s.getJobStats(); + long failed = js.getFailedRecords() != null ? js.getFailedRecords() : 0; + long success = js.getSuccessRecords() != null ? js.getSuccessRecords() : 0; + long total = js.getTotalRecords() != null ? js.getTotalRecords() : 0; + return failed > 0 || (total > 0 && success < total); + } + + private Stats initializeStats(ReindexingConfiguration config, Set entities) { + Stats s = new Stats(); + s.setEntityStats(new org.openmetadata.schema.system.EntityStats()); + s.setJobStats(new StepStats()); + s.setReaderStats(new StepStats()); + s.setSinkStats(new StepStats()); + + int total = 0; + for (String entityType : entities) { + int entityTotal = getEntityTotal(entityType, config); + total += entityTotal; + StepStats es = new StepStats(); + es.setTotalRecords(entityTotal); + es.setSuccessRecords(0); + es.setFailedRecords(0); + s.getEntityStats().getAdditionalProperties().put(entityType, es); + } + + if (entities.contains(Entity.TABLE) && !entities.contains(Entity.TABLE_COLUMN)) { + StepStats columnStats = new StepStats(); + columnStats.setTotalRecords(0); + columnStats.setSuccessRecords(0); + columnStats.setFailedRecords(0); + s.getEntityStats().getAdditionalProperties().put(Entity.TABLE_COLUMN, columnStats); + } + + s.getJobStats().setTotalRecords(total); + s.getJobStats().setSuccessRecords(0); + s.getJobStats().setFailedRecords(0); + s.getReaderStats().setTotalRecords(total); + s.getReaderStats().setSuccessRecords(0); + s.getReaderStats().setFailedRecords(0); + s.getReaderStats().setWarningRecords(0); + s.getSinkStats().setTotalRecords(0); + s.getSinkStats().setSuccessRecords(0); + s.getSinkStats().setFailedRecords(0); + + s.setProcessStats(new StepStats()); + s.getProcessStats().setTotalRecords(0); + s.getProcessStats().setSuccessRecords(0); + s.getProcessStats().setFailedRecords(0); + return s; + } + + private int getEntityTotal(String entityType, ReindexingConfiguration config) { + try { + if (!EntityReader.TIME_SERIES_ENTITIES.contains(entityType)) { + EntityRepository repository = Entity.getEntityRepository(entityType); + return repository + .getDao() + .listCount(new ListFilter(org.openmetadata.schema.type.Include.ALL)); + } + + EntityTimeSeriesRepository repository; + ListFilter listFilter = new ListFilter(null); + if (isDataInsightIndex(entityType)) { + listFilter.addQueryParam("entityFQNHash", FullyQualifiedName.buildHash(entityType)); + repository = Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA); + } else { + repository = Entity.getEntityTimeSeriesRepository(entityType); + } + + long startTs = config != null ? config.getTimeSeriesStartTs(entityType) : -1; + if (startTs > 0) { + long endTs = System.currentTimeMillis(); + return repository.getTimeSeriesDao().listCount(listFilter, startTs, endTs, false); + } + return repository.getTimeSeriesDao().listCount(listFilter); + } catch (Exception e) { + LOG.debug("Error getting total records for '{}'", entityType, e); + return 0; + } + } + + private int getTotalEntityRecords(String entityType) { + StepStats es = + stats.get() != null + && stats.get().getEntityStats() != null + && stats.get().getEntityStats().getAdditionalProperties() != null + ? stats.get().getEntityStats().getAdditionalProperties().get(entityType) + : null; + if (es != null && es.getTotalRecords() != null) { + return es.getTotalRecords(); + } + return 0; + } + + private synchronized void updateReaderStats(int success, int failed, int warnings) { + Stats s = stats.get(); + if (s == null) return; + StepStats rs = s.getReaderStats(); + if (rs == null) { + rs = new StepStats(); + s.setReaderStats(rs); + } + rs.setSuccessRecords((rs.getSuccessRecords() != null ? rs.getSuccessRecords() : 0) + success); + rs.setFailedRecords((rs.getFailedRecords() != null ? rs.getFailedRecords() : 0) + failed); + rs.setWarningRecords((rs.getWarningRecords() != null ? rs.getWarningRecords() : 0) + warnings); + } + + private synchronized void updateEntityAndJobStats(String entityType, StepStats entityDelta) { + Stats s = stats.get(); + if (s == null || s.getEntityStats() == null) return; + + StepStats es = s.getEntityStats().getAdditionalProperties().get(entityType); + if (es != null) { + es.setSuccessRecords(es.getSuccessRecords() + entityDelta.getSuccessRecords()); + es.setFailedRecords(es.getFailedRecords() + entityDelta.getFailedRecords()); + } + + StepStats js = s.getJobStats(); + if (js != null) { + int totalSuccess = + s.getEntityStats().getAdditionalProperties().entrySet().stream() + .filter(e -> !Entity.TABLE_COLUMN.equals(e.getKey())) + .mapToInt(e -> e.getValue().getSuccessRecords()) + .sum(); + int totalFailed = + s.getEntityStats().getAdditionalProperties().entrySet().stream() + .filter(e -> !Entity.TABLE_COLUMN.equals(e.getKey())) + .mapToInt(e -> e.getValue().getFailedRecords()) + .sum(); + js.setSuccessRecords(totalSuccess); + js.setFailedRecords(totalFailed); + } + } + + private synchronized void syncSinkStats() { + if (searchIndexSink == null) return; + Stats s = stats.get(); + if (s == null) return; + + StepStats bulkStats = searchIndexSink.getStats(); + if (bulkStats == null) return; + + StepStats sinkStats = s.getSinkStats(); + if (sinkStats == null) { + sinkStats = new StepStats(); + s.setSinkStats(sinkStats); + } + sinkStats.setTotalRecords( + bulkStats.getTotalRecords() != null ? bulkStats.getTotalRecords() : 0); + sinkStats.setSuccessRecords( + bulkStats.getSuccessRecords() != null ? bulkStats.getSuccessRecords() : 0); + sinkStats.setFailedRecords( + bulkStats.getFailedRecords() != null ? bulkStats.getFailedRecords() : 0); + + StepStats vectorStats = searchIndexSink.getVectorStats(); + if (vectorStats != null + && vectorStats.getTotalRecords() != null + && vectorStats.getTotalRecords() > 0) { + s.setVectorStats(vectorStats); + } + + StepStats processStats = searchIndexSink.getProcessStats(); + if (processStats != null) { + s.setProcessStats(processStats); + } + } + + private void updateColumnStatsFromSink() { + if (searchIndexSink == null) return; + Stats s = stats.get(); + if (s == null || s.getEntityStats() == null) return; + + StepStats columnStats = searchIndexSink.getColumnStats(); + if (columnStats != null && columnStats.getTotalRecords() > 0) { + StepStats existing = s.getEntityStats().getAdditionalProperties().get(Entity.TABLE_COLUMN); + if (existing != null) { + existing.setTotalRecords(columnStats.getTotalRecords()); + existing.setSuccessRecords(columnStats.getSuccessRecords()); + existing.setFailedRecords(columnStats.getFailedRecords()); + } + } + } + + private void shutdownExecutors() { + shutdownExecutor(producerExecutor, "producer"); + shutdownExecutor(jobExecutor, "job"); + shutdownExecutor(consumerExecutor, "consumer"); + } + + private void shutdownExecutor(ExecutorService executor, String name) { + if (executor != null && !executor.isShutdown()) { + executor.shutdown(); + try { + if (!executor.awaitTermination(30, TimeUnit.SECONDS)) { + executor.shutdownNow(); + LOG.warn("{} executor did not terminate in time", name); + } + } catch (InterruptedException e) { + executor.shutdownNow(); + Thread.currentThread().interrupt(); + } + } + } + + public void stop() { + stopped.set(true); + if (entityReader != null) entityReader.stop(); + + if (searchIndexSink != null) { + LOG.info( + "Stopping pipeline: flushing sink ({} active bulk requests)", + searchIndexSink.getActiveBulkRequestCount()); + searchIndexSink.flushAndAwait(10); + } + + int dropped = taskQueue != null ? taskQueue.size() : 0; + if (dropped > 0) { + LOG.warn("Dropping {} queued tasks during shutdown", dropped); + } + + if (taskQueue != null) { + taskQueue.clear(); + for (int i = 0; i < MAX_CONSUMER_THREADS; i++) { + taskQueue.offer(new IndexingTask<>(POISON_PILL, null, -1)); + } + } + shutdownExecutors(); + } + + @Override + public void close() { + stop(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingStrategy.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingStrategy.java new file mode 100644 index 000000000000..e7d4b2018b9e --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingStrategy.java @@ -0,0 +1,21 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Optional; +import org.openmetadata.schema.system.Stats; + +/** + * Strategy interface for reindexing execution. Encapsulates the differences between single-server + * and distributed indexing so that SearchIndexApp uses a single code path regardless of mode. + */ +public interface IndexingStrategy { + + void addListener(ReindexingProgressListener listener); + + ExecutionResult execute(ReindexingConfiguration config, ReindexingJobContext context); + + Optional getStats(); + + void stop(); + + boolean isStopped(); +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java index 6b290b3627c0..1fb84174d3a4 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrchestratorContext.java @@ -28,5 +28,5 @@ public interface OrchestratorContext { ReindexingProgressListener createProgressListener(EventPublisherJob jobData); - ReindexingJobContext createReindexingContext(); + ReindexingJobContext createReindexingContext(boolean distributed); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrphanedIndexCleaner.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrphanedIndexCleaner.java index 94dc0ca0aedc..3f4c85b3f97a 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrphanedIndexCleaner.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/OrphanedIndexCleaner.java @@ -25,7 +25,7 @@ * considered orphaned if: * *
    - *
  • It contains "_rebuild_" in its name (created during staged reindexing) + *
  • It contains "_rebuild_" in its name (created during recreateIndex=true) *
  • It has ZERO aliases pointing to it (not serving any traffic) *
* diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContext.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContext.java index b9ea0bbf4e72..e39a619749b8 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContext.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContext.java @@ -14,8 +14,9 @@ public class QuartzJobContext implements ReindexingJobContext { private final String jobName; private final Long startTime; private final UUID appId; + private final boolean distributed; - public QuartzJobContext(JobExecutionContext jobExecutionContext, App app) { + public QuartzJobContext(JobExecutionContext jobExecutionContext, App app, boolean distributed) { this.jobName = jobExecutionContext != null ? jobExecutionContext.getJobDetail().getKey().getName() @@ -23,6 +24,7 @@ public QuartzJobContext(JobExecutionContext jobExecutionContext, App app) { this.startTime = System.currentTimeMillis(); this.appId = app != null ? app.getId() : null; this.jobId = appId != null ? appId : UUID.randomUUID(); + this.distributed = distributed; } @Override @@ -45,6 +47,11 @@ public UUID getAppId() { return appId; } + @Override + public boolean isDistributed() { + return distributed; + } + @Override public String getSource() { return "QUARTZ"; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java index 379f80fe02cf..497616eac5e1 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContext.java @@ -92,7 +92,7 @@ public ReindexingProgressListener createProgressListener(EventPublisherJob jobDa } @Override - public ReindexingJobContext createReindexingContext() { - return new QuartzJobContext(ctx, app); + public ReindexingJobContext createReindexingContext(boolean distributed) { + return new QuartzJobContext(ctx, app, distributed); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java index 7d7a68357156..2426e63c3672 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfiguration.java @@ -25,7 +25,9 @@ public record ReindexingConfiguration( int fieldFetchThreads, int docBuildThreads, long statsIntervalMs, + boolean recreateIndex, boolean autoTune, + boolean useDistributedIndexing, boolean force, int maxRetries, int initialBackoff, @@ -75,7 +77,9 @@ public static ReindexingConfiguration applyAutoTuning( .fieldFetchThreads(metrics.getRecommendedFieldFetchThreads()) .docBuildThreads(metrics.getRecommendedDocBuildThreads()) .statsIntervalMs(metrics.getRecommendedStatsIntervalMs()) + .recreateIndex(config.recreateIndex()) .autoTune(true) + .useDistributedIndexing(config.useDistributedIndexing()) .force(config.force()) .maxRetries(config.maxRetries()) .initialBackoff(config.initialBackoff()) @@ -124,7 +128,9 @@ public static ReindexingConfiguration from(EventPublisherJob jobData) { DEFAULT_FIELD_FETCH_THREADS, DEFAULT_DOC_BUILD_THREADS, DEFAULT_STATS_INTERVAL_MS, + Boolean.TRUE.equals(jobData.getRecreateIndex()), Boolean.TRUE.equals(jobData.getAutoTune()), + Boolean.TRUE.equals(jobData.getUseDistributedIndexing()), Boolean.TRUE.equals(jobData.getForce()), jobData.getMaxRetries() != null ? jobData.getMaxRetries() : DEFAULT_MAX_RETRIES, jobData.getInitialBackoff() != null ? jobData.getInitialBackoff() : DEFAULT_INITIAL_BACKOFF, @@ -181,9 +187,7 @@ public boolean hasSlackConfig() { /** Check if this is a subset (smart) reindexing */ public boolean isSmartReindexing() { - return entities != null - && !entities.contains(SearchIndexEntityTypes.ALL) - && entities.size() < 20; + return entities != null && !entities.contains("all") && entities.size() < 20 && recreateIndex; } /** Creates a builder for more flexible configuration creation */ @@ -202,7 +206,9 @@ public static class Builder { private int fieldFetchThreads = DEFAULT_FIELD_FETCH_THREADS; private int docBuildThreads = DEFAULT_DOC_BUILD_THREADS; private long statsIntervalMs = DEFAULT_STATS_INTERVAL_MS; + private boolean recreateIndex = false; private boolean autoTune = false; + private boolean useDistributedIndexing = false; private boolean force = false; private int maxRetries = DEFAULT_MAX_RETRIES; private int initialBackoff = DEFAULT_INITIAL_BACKOFF; @@ -264,11 +270,21 @@ public Builder statsIntervalMs(long statsIntervalMs) { return this; } + public Builder recreateIndex(boolean recreateIndex) { + this.recreateIndex = recreateIndex; + return this; + } + public Builder autoTune(boolean autoTune) { this.autoTune = autoTune; return this; } + public Builder useDistributedIndexing(boolean useDistributedIndexing) { + this.useDistributedIndexing = useDistributedIndexing; + return this; + } + public Builder force(boolean force) { this.force = force; return this; @@ -331,7 +347,9 @@ public ReindexingConfiguration build() { fieldFetchThreads, docBuildThreads, statsIntervalMs, + recreateIndex, autoTune, + useDistributedIndexing, force, maxRetries, initialBackoff, diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingJobContext.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingJobContext.java index a5a9b6679e86..8b1b6af73231 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingJobContext.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingJobContext.java @@ -20,6 +20,9 @@ public interface ReindexingJobContext { /** Application ID (for Quartz-based jobs, null for CLI/API) */ UUID getAppId(); + /** Whether this is a distributed indexing job */ + boolean isDistributed(); + /** The source that triggered this job (e.g., "QUARTZ", "CLI", "API") */ String getSource(); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java index ca65907652c2..05c9e1f4ae1c 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestrator.java @@ -33,13 +33,14 @@ @Slf4j public class ReindexingOrchestrator { + private static final String ALL = "all"; private final CollectionDAO collectionDAO; private final SearchRepository searchRepository; private final OrchestratorContext context; @Getter private EventPublisherJob jobData; private volatile boolean stopped = false; - private volatile DistributedIndexingStrategy activeStrategy; + private volatile IndexingStrategy activeStrategy; private volatile Map resultMetadata = Collections.emptyMap(); public ReindexingOrchestrator( @@ -94,7 +95,7 @@ public void stop() { LOG.info("Reindexing job is being stopped."); stopped = true; - DistributedIndexingStrategy strategy = this.activeStrategy; + IndexingStrategy strategy = this.activeStrategy; if (strategy != null) { try { strategy.stop(); @@ -109,7 +110,6 @@ public void stop() { AppRunRecord appRecord = context.getJobRecord(); appRecord.setStatus(AppRunRecord.Status.STOPPED); - sanitizeRunRecordConfig(appRecord); OmAppJobListener.fillTerminalTimings(appRecord); context.storeRunRecord(JsonUtils.pojoToJson(appRecord)); context.pushStatusUpdate(appRecord, true); @@ -129,10 +129,10 @@ private void initializeJobData() { jobData = loadJobData(); } - if (ON_DEMAND_JOB.equals(context.getJobName())) { + String jobName = context.getJobName(); + if (jobName.equals(ON_DEMAND_JOB)) { Map jsonAppConfig = JsonUtils.convertValue(jobData, new TypeReference>() {}); - SearchIndexAppConfigSanitizer.removeRemovedOptions(jsonAppConfig); context.updateAppConfiguration(jsonAppConfig); } } @@ -140,18 +140,12 @@ private void initializeJobData() { private EventPublisherJob loadJobData() { String appConfigJson = context.getAppConfigJson(); if (appConfigJson != null) { - Map appConfig = - JsonUtils.readValue(appConfigJson, new TypeReference>() {}); - return JsonUtils.convertValue( - SearchIndexAppConfigSanitizer.copyWithoutRemovedOptions(appConfig), - EventPublisherJob.class); + return JsonUtils.readValue(appConfigJson, EventPublisherJob.class); } Map appConfig = context.getAppConfiguration(); if (appConfig != null) { - return JsonUtils.convertValue( - SearchIndexAppConfigSanitizer.copyWithoutRemovedOptions(appConfig), - EventPublisherJob.class); + return JsonUtils.convertValue(appConfig, EventPublisherJob.class); } LOG.error("Unable to initialize jobData from JobDataMap or App configuration"); @@ -218,77 +212,44 @@ private void cleanupOrphanedIndicesPreFlight() { } private void runReindexing() { - if (hasNoEntitiesSelected()) { - completeWithoutEntities(); + if (jobData.getEntities() == null || jobData.getEntities().isEmpty()) { + LOG.info("No entities selected for reindexing, completing immediately"); + jobData.setStatus(EventPublisherJob.Status.COMPLETED); + jobData.setStats(new Stats()); return; } setupEntities(); cleanupOldFailures(); - logJobStart(); - DistributedIndexingStrategy strategy = createDistributedStrategy(); - activeStrategy = strategy; - registerProgressListeners(strategy); - - ReindexingConfiguration config = buildReindexingConfiguration(); - ExecutionResult result = executeDistributedReindex(strategy, config); - persistExecutionResult(result); - } - - private boolean hasNoEntitiesSelected() { - return jobData.getEntities() == null || jobData.getEntities().isEmpty(); - } - - private void completeWithoutEntities() { - LOG.info("No entities selected for reindexing, completing immediately"); - jobData.setStatus(EventPublisherJob.Status.COMPLETED); - jobData.setStats(new Stats()); - } - - private void logJobStart() { LOG.info( - "Search Index Job Started for Entities: {} using staged index promotion", - jobData.getEntities()); - } + "Search Index Job Started for Entities: {}, RecreateIndex: {}, DistributedIndexing: {}", + jobData.getEntities(), + jobData.getRecreateIndex(), + jobData.getUseDistributedIndexing()); - private DistributedIndexingStrategy createDistributedStrategy() { - AppRunRecord appRecord = context.getJobRecord(); - return new DistributedIndexingStrategy( - collectionDAO, - searchRepository, - jobData, - appRecord.getAppId(), - appRecord.getStartTime(), - context.getJobName()); - } + activeStrategy = createStrategy(); - private void registerProgressListeners(DistributedIndexingStrategy strategy) { - strategy.addListener(context.createProgressListener(jobData)); - strategy.addListener(new LoggingProgressListener()); + activeStrategy.addListener(context.createProgressListener(jobData)); + activeStrategy.addListener(new LoggingProgressListener()); if (hasSlackConfig()) { - strategy.addListener( + String instanceUrl = getInstanceUrl(); + activeStrategy.addListener( new SlackProgressListener( - jobData.getSlackBotToken(), jobData.getSlackChannel(), getInstanceUrl())); + jobData.getSlackBotToken(), jobData.getSlackChannel(), instanceUrl)); } - } - private ReindexingConfiguration buildReindexingConfiguration() { + ReindexingJobContext jobContext = + context.createReindexingContext(Boolean.TRUE.equals(jobData.getUseDistributedIndexing())); + ReindexingConfiguration config = ReindexingConfiguration.from(jobData); - config = - ReindexingConfiguration.applyAutoTuning(config, searchRepository, countTotalEntities()); + long totalEntities = countTotalEntities(); + config = ReindexingConfiguration.applyAutoTuning(config, searchRepository, totalEntities); config.applyTo(jobData); updateRunRecordConfig(config); - return config; - } - - private ExecutionResult executeDistributedReindex( - DistributedIndexingStrategy strategy, ReindexingConfiguration config) { - return strategy.execute(config, context.createReindexingContext()); - } - private void persistExecutionResult(ExecutionResult result) { + ExecutionResult result = activeStrategy.execute(config, jobContext); updateJobDataFromResult(result); if (jobData.getStats() != null) { @@ -300,6 +261,20 @@ private void persistExecutionResult(ExecutionResult result) { } } + private IndexingStrategy createStrategy() { + if (Boolean.TRUE.equals(jobData.getUseDistributedIndexing())) { + AppRunRecord appRecord = context.getJobRecord(); + return new DistributedIndexingStrategy( + collectionDAO, + searchRepository, + jobData, + appRecord.getAppId(), + appRecord.getStartTime(), + context.getJobName()); + } + return new SingleServerIndexingStrategy(collectionDAO, searchRepository); + } + private void updateJobDataFromResult(ExecutionResult result) { if (result.finalStats() != null) { Stats stats = result.finalStats(); @@ -323,7 +298,6 @@ private void updateRunRecordConfig(ReindexingConfiguration config) { if (appRecord != null) { Map configMap = appRecord.getConfig(); if (configMap != null) { - SearchIndexAppConfigSanitizer.removeRemovedOptions(configMap); configMap.put("batchSize", config.batchSize()); configMap.put("consumerThreads", config.consumerThreads()); configMap.put("producerThreads", config.producerThreads()); @@ -362,7 +336,7 @@ private void saveResultMetadataToJobRecord(Map metadata) { } private void handleExecutionException(Exception ex) { - DistributedIndexingStrategy strategy = this.activeStrategy; + IndexingStrategy strategy = this.activeStrategy; if (strategy != null && jobData != null) { try { strategy.getStats().ifPresent(jobData::setStats); @@ -395,7 +369,6 @@ private void finalizeJobExecution() { if (stopped) { AppRunRecord appRecord = context.getJobRecord(); appRecord.setStatus(AppRunRecord.Status.STOPPED); - sanitizeRunRecordConfig(appRecord); OmAppJobListener.fillTerminalTimings(appRecord); context.storeRunRecord(JsonUtils.pojoToJson(appRecord)); } @@ -412,7 +385,6 @@ private void sendUpdates() { private void updateRecordToDbAndNotify() { AppRunRecord appRecord = context.getJobRecord(); appRecord.setStatus(AppRunRecord.Status.fromValue(jobData.getStatus().value())); - sanitizeRunRecordConfig(appRecord); OmAppJobListener.fillTerminalTimings(appRecord); if (jobData.getFailure() != null) { @@ -465,12 +437,6 @@ private void updateRecordToDbAndNotify() { } } - private void sanitizeRunRecordConfig(AppRunRecord appRecord) { - if (appRecord != null) { - SearchIndexAppConfigSanitizer.removeRemovedOptions(appRecord.getConfig()); - } - } - private void cleanupOldFailures() { try { int deleted = collectionDAO.searchIndexFailureDAO().deleteAll(); @@ -500,11 +466,10 @@ private void cleanupOrphanedIndices() { } private void setupEntities() { - Set entities = - jobData.getEntities().contains(SearchIndexEntityTypes.ALL) - ? getAll() - : jobData.getEntities(); - jobData.setEntities(SearchIndexEntityTypes.normalizeEntityTypes(entities)); + boolean containsAll = jobData.getEntities().contains(ALL); + if (containsAll) { + jobData.setEntities(getAll()); + } } private Set getAll() { @@ -522,10 +487,9 @@ private long countTotalEntities() { long total = 0; for (String entityType : jobData.getEntities()) { try { - String normalizedEntityType = SearchIndexEntityTypes.normalizeEntityType(entityType); - if (!SearchIndexEntityTypes.isTimeSeriesEntity(normalizedEntityType)) { + if (!SearchIndexApp.TIME_SERIES_ENTITIES.contains(entityType)) { total += - Entity.getEntityRepository(normalizedEntityType) + Entity.getEntityRepository(entityType) .getDao() .listCount( new org.openmetadata.service.jdbi3.ListFilter( diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java index 99846826ce42..d04b613a3ac4 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingProgressListener.java @@ -33,7 +33,7 @@ default void onJobStarted(ReindexingJobContext context) {} /** Called when job configuration is determined (after auto-tune) */ default void onJobConfigured(ReindexingJobContext context, ReindexingConfiguration config) {} - /** Called when staged index preparation begins. */ + /** Called when index recreation begins (if recreateIndex=true) */ default void onIndexRecreationStarted(Set entities) {} /** Called when a specific entity type processing begins */ diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java index 7122574a3c97..3f4b1c93b602 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexApp.java @@ -1,17 +1,22 @@ package org.openmetadata.service.apps.bundles.searchIndex; +import static org.openmetadata.service.Entity.QUERY_COST_RECORD; +import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; +import static org.openmetadata.service.Entity.TEST_CASE_RESULT; + import jakarta.ws.rs.core.Response; import java.util.List; import java.util.Map; +import java.util.Set; import lombok.Getter; import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.analytics.ReportData; import org.openmetadata.schema.entity.app.App; import org.openmetadata.schema.entity.app.AppRunRecord; import org.openmetadata.schema.system.EventPublisherJob; import org.openmetadata.schema.utils.JsonUtils; import org.openmetadata.service.apps.AbstractNativeApplication; import org.openmetadata.service.apps.bundles.searchIndex.distributed.DistributedSearchIndexCoordinator; -import org.openmetadata.service.apps.bundles.searchIndex.distributed.IndexJobStatus; import org.openmetadata.service.exception.AppException; import org.openmetadata.service.jdbi3.AppRepository; import org.openmetadata.service.jdbi3.CollectionDAO; @@ -20,12 +25,6 @@ @Slf4j public class SearchIndexApp extends AbstractNativeApplication { - private static final String REINDEX_LOCK_KEY = "SEARCH_REINDEX_LOCK"; - private static final List ACTIVE_DISTRIBUTED_JOB_STATUSES = - List.of( - IndexJobStatus.RUNNING.name(), - IndexJobStatus.READY.name(), - IndexJobStatus.INITIALIZING.name()); public static class ReindexingException extends RuntimeException { public ReindexingException(String message) { @@ -37,6 +36,17 @@ public ReindexingException(String message, Throwable cause) { } } + public static final Set TIME_SERIES_ENTITIES = + Set.of( + ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), + ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), + ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), + TEST_CASE_RESOLUTION_STATUS, + TEST_CASE_RESULT, + QUERY_COST_RECORD); + @Getter private EventPublisherJob jobData; private volatile ReindexingOrchestrator orchestrator; @@ -47,10 +57,7 @@ public SearchIndexApp(CollectionDAO collectionDAO, SearchRepository searchReposi @Override public void init(App app) { super.init(app); - Map appConfig = - SearchIndexAppConfigSanitizer.copyWithoutRemovedOptions( - JsonUtils.getMap(app.getAppConfiguration())); - jobData = JsonUtils.convertValue(appConfig, EventPublisherJob.class); + jobData = JsonUtils.convertValue(app.getAppConfiguration(), EventPublisherJob.class); } @Override @@ -108,7 +115,6 @@ private void updateRunRecordToStopped() { run -> { run.withStatus(AppRunRecord.Status.STOPPED); run.withEndTime(System.currentTimeMillis()); - SearchIndexAppConfigSanitizer.removeRemovedOptions(run.getConfig()); appRepository.updateAppStatus(app.getId(), run); LOG.info("Updated app run record to STOPPED for {}", app.getName()); }); @@ -126,7 +132,9 @@ public void uninstall() { private void purgeSearchIndexTables() { List activeJobs = - collectionDAO.searchIndexJobDAO().findByStatuses(ACTIVE_DISTRIBUTED_JOB_STATUSES); + collectionDAO + .searchIndexJobDAO() + .findByStatuses(List.of("RUNNING", "READY", "INITIALIZING")); if (!activeJobs.isEmpty()) { LOG.warn( "Uninstalling SearchIndexApp while {} distributed job(s) are still active. " @@ -139,7 +147,7 @@ private void purgeSearchIndexTables() { .searchIndexJobDAO() .update( job.id(), - IndexJobStatus.STOPPED.name(), + "STOPPED", job.processedRecords(), job.successRecords(), job.failedRecords(), @@ -158,7 +166,7 @@ private void purgeSearchIndexTables() { () -> collectionDAO.searchIndexPartitionDAO().deleteAll(), () -> collectionDAO.searchIndexServerStatsDAO().deleteAll(), () -> collectionDAO.searchIndexFailureDAO().deleteAll(), - () -> collectionDAO.searchReindexLockDAO().delete(REINDEX_LOCK_KEY), + () -> collectionDAO.searchReindexLockDAO().delete("SEARCH_REINDEX_LOCK"), () -> collectionDAO.searchIndexJobDAO().deleteAll(), () -> { App app = getApp(); @@ -177,9 +185,7 @@ private void purgeSearchIndexTables() { @Override protected void validateConfig(Map appConfig) { try { - JsonUtils.convertValue( - SearchIndexAppConfigSanitizer.copyWithoutRemovedOptions(appConfig), - EventPublisherJob.class); + JsonUtils.convertValue(appConfig, EventPublisherJob.class); } catch (IllegalArgumentException e) { throw AppException.byMessage( Response.Status.BAD_REQUEST, "Invalid App Configuration: " + e.getMessage()); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppConfigSanitizer.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppConfigSanitizer.java deleted file mode 100644 index 054edc1b65a7..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppConfigSanitizer.java +++ /dev/null @@ -1,28 +0,0 @@ -package org.openmetadata.service.apps.bundles.searchIndex; - -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.Set; - -final class SearchIndexAppConfigSanitizer { - private static final Set REMOVED_OPTIONS = - Set.of("recreateIndex", "useDistributedIndexing"); - - private SearchIndexAppConfigSanitizer() {} - - static Map copyWithoutRemovedOptions(Map config) { - if (config == null) { - return config; - } - Map sanitized = new LinkedHashMap<>(config); - removeRemovedOptions(sanitized); - return sanitized; - } - - static void removeRemovedOptions(Map config) { - if (config == null || config.isEmpty()) { - return; - } - REMOVED_OPTIONS.forEach(config::remove); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexEntityTypes.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexEntityTypes.java deleted file mode 100644 index c83329a77bee..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexEntityTypes.java +++ /dev/null @@ -1,63 +0,0 @@ -/* - * Copyright 2024 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.service.apps.bundles.searchIndex; - -import static org.openmetadata.service.Entity.QUERY_COST_RECORD; -import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; -import static org.openmetadata.service.Entity.TEST_CASE_RESULT; - -import java.util.LinkedHashSet; -import java.util.Set; -import org.openmetadata.schema.analytics.ReportData; - -public final class SearchIndexEntityTypes { - public static final String ALL = "all"; - public static final String QUERY_COST_RESULT = "queryCostResult"; - - public static final Set TIME_SERIES_ENTITIES = - Set.of( - ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), - ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), - ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), - ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), - ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), - TEST_CASE_RESOLUTION_STATUS, - TEST_CASE_RESULT, - QUERY_COST_RECORD); - - private SearchIndexEntityTypes() {} - - public static String normalizeEntityType(String entityType) { - return QUERY_COST_RESULT.equals(entityType) ? QUERY_COST_RECORD : entityType; - } - - public static Set normalizeEntityTypes(Set entityTypes) { - if (entityTypes == null || entityTypes.isEmpty()) { - return entityTypes; - } - Set normalizedEntityTypes = new LinkedHashSet<>(); - for (String entityType : entityTypes) { - normalizedEntityTypes.add(normalizeEntityType(entityType)); - } - return normalizedEntityTypes; - } - - public static boolean isTimeSeriesEntity(String entityType) { - return TIME_SERIES_ENTITIES.contains(normalizeEntityType(entityType)); - } - - public static boolean isDataInsightEntity(String entityType) { - return entityType != null && entityType.endsWith("ReportData"); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutor.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutor.java new file mode 100644 index 000000000000..60591421d4c7 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutor.java @@ -0,0 +1,1950 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; +import static org.openmetadata.common.utils.CommonUtil.nullOrEmpty; +import static org.openmetadata.service.Entity.QUERY_COST_RECORD; +import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; +import static org.openmetadata.service.Entity.TEST_CASE_RESULT; +import static org.openmetadata.service.workflows.searchIndex.ReindexingUtil.RECREATE_CONTEXT; +import static org.openmetadata.service.workflows.searchIndex.ReindexingUtil.TARGET_INDEX_KEY; +import static org.openmetadata.service.workflows.searchIndex.ReindexingUtil.isDataInsightIndex; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Phaser; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.exception.ExceptionUtils; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.EntityTimeSeriesInterface; +import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.system.EntityStats; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.schema.type.Include; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.stats.EntityStatsTracker; +import org.openmetadata.service.apps.bundles.searchIndex.stats.JobStatsManager; +import org.openmetadata.service.apps.bundles.searchIndex.stats.StageStatsTracker; +import org.openmetadata.service.exception.SearchIndexException; +import org.openmetadata.service.jdbi3.BoundedListFilter; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.EntityRepository; +import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; +import org.openmetadata.service.jdbi3.ListFilter; +import org.openmetadata.service.search.DefaultRecreateHandler; +import org.openmetadata.service.search.EntityReindexContext; +import org.openmetadata.service.search.RecreateIndexHandler; +import org.openmetadata.service.search.ReindexContext; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.util.FullyQualifiedName; +import org.openmetadata.service.util.RestUtil; +import org.openmetadata.service.workflows.interfaces.Source; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntitiesSource; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntityTimeSeriesSource; +import org.slf4j.MDC; + +/** + * Core reindexing executor that handles entity indexing without any Quartz dependencies. Can be + * used by: + * + *
    + *
  • SearchIndexApp (Quartz integration) + *
  • CLI tools + *
  • REST API endpoints + *
  • Unit tests + *
+ * + *

Uses ReindexingProgressListener for extensible progress reporting. + */ +@Slf4j +public class SearchIndexExecutor implements AutoCloseable { + + private static final String ALL = "all"; + private static final String POISON_PILL = "__POISON_PILL__"; + private static final int DEFAULT_QUEUE_SIZE = 20000; + private static final String RECREATE_INDEX = "recreateIndex"; + private static final String ENTITY_TYPE_KEY = "entityType"; + private static final String QUERY_COST_RESULT_INCORRECT = "queryCostResult"; + private static final String QUERY_COST_RESULT_WARNING = + "Found incorrect entity type 'queryCostResult', correcting to 'queryCostRecord'"; + + private static final int AVAILABLE_PROCESSORS = Runtime.getRuntime().availableProcessors(); + private static final int MAX_READERS_PER_ENTITY = 5; + private static final int MAX_PRODUCER_THREADS = Math.min(20, AVAILABLE_PROCESSORS * 2); + private static final int MAX_CONSUMER_THREADS = Math.min(20, AVAILABLE_PROCESSORS * 2); + private static final int MAX_TOTAL_THREADS = Math.min(50, AVAILABLE_PROCESSORS * 4); + + public static final Set TIME_SERIES_ENTITIES = + Set.of( + ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), + ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), + ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), + TEST_CASE_RESOLUTION_STATUS, + TEST_CASE_RESULT, + QUERY_COST_RECORD); + + private final CollectionDAO collectionDAO; + private final SearchRepository searchRepository; + private final CompositeProgressListener listeners; + private final AtomicBoolean stopped = new AtomicBoolean(false); + private final AtomicBoolean sinkClosed = new AtomicBoolean(false); + + private BulkSink searchIndexSink; + private RecreateIndexHandler recreateIndexHandler; + private ReindexContext recreateContext; + private ExecutorService producerExecutor; + private ExecutorService consumerExecutor; + private ExecutorService jobExecutor; + private BlockingQueue> taskQueue; + private final AtomicBoolean producersDone = new AtomicBoolean(false); + + @Getter private final AtomicReference stats = new AtomicReference<>(); + private final AtomicReference batchSize = new AtomicReference<>(100); + + private ReindexingConfiguration config; + private ReindexingJobContext context; + private long startTime; + private IndexingFailureRecorder failureRecorder; + private JobStatsManager statsManager; + private final Map entityBatchCounters = new ConcurrentHashMap<>(); + private final Map entityBatchFailures = new ConcurrentHashMap<>(); + private final Set promotedEntities = ConcurrentHashMap.newKeySet(); + private final Map sinkTrackers = new ConcurrentHashMap<>(); + private final Map> contextDataCache = new ConcurrentHashMap<>(); + private static final long SINK_SYNC_INTERVAL_MS = 2000; + private final AtomicLong lastSinkSyncTime = new AtomicLong(0); + + record IndexingTask(String entityType, ResultList entities, int offset, int retryCount) { + IndexingTask(String entityType, ResultList entities, int offset) { + this(entityType, entities, offset, 0); + } + } + + record ThreadConfiguration(int numProducers, int numConsumers) {} + + @FunctionalInterface + interface KeysetBatchReader { + ResultList readNextKeyset(String cursor) throws SearchIndexException; + } + + static class MemoryInfo { + final long maxMemory; + final long usedMemory; + final double usageRatio; + + MemoryInfo() { + Runtime runtime = Runtime.getRuntime(); + this.maxMemory = runtime.maxMemory(); + long totalMemory = runtime.totalMemory(); + long freeMemory = runtime.freeMemory(); + this.usedMemory = totalMemory - freeMemory; + this.usageRatio = (double) usedMemory / maxMemory; + } + } + + public SearchIndexExecutor(CollectionDAO collectionDAO, SearchRepository searchRepository) { + this.collectionDAO = collectionDAO; + this.searchRepository = searchRepository; + this.listeners = new CompositeProgressListener(); + } + + private EntityStatsTracker getTracker(String entityType) { + return statsManager != null ? statsManager.getTracker(entityType) : null; + } + + private void initStatsManager() { + if (statsManager == null && context != null) { + String jobId = context.getJobId().toString(); + String serverId = + org.openmetadata + .service + .apps + .bundles + .searchIndex + .distributed + .ServerIdentityResolver + .getInstance() + .getServerId(); + statsManager = new JobStatsManager(jobId, serverId, collectionDAO); + } + } + + public SearchIndexExecutor addListener(ReindexingProgressListener listener) { + listeners.addListener(listener); + return this; + } + + public SearchIndexExecutor removeListener(ReindexingProgressListener listener) { + listeners.removeListener(listener); + return this; + } + + /** + * Execute reindexing with the given configuration. + * + * @param config The reindexing configuration + * @param context The job context + * @return ExecutionResult with final stats + */ + public ExecutionResult execute(ReindexingConfiguration config, ReindexingJobContext context) { + this.config = config; + this.context = context; + this.startTime = System.currentTimeMillis(); + initializeState(); + + listeners.onJobStarted(context); + + try { + return executeSingleServer(); + } catch (Exception e) { + LOG.error("Reindexing failed", e); + listeners.onJobFailed(stats.get(), e); + return ExecutionResult.fromStats(stats.get(), ExecutionResult.Status.FAILED, startTime); + } + } + + private void initializeState() { + stopped.set(false); + sinkClosed.set(false); + recreateContext = null; + producersDone.set(false); + entityBatchCounters.clear(); + entityBatchFailures.clear(); + promotedEntities.clear(); + sinkTrackers.clear(); + contextDataCache.clear(); + lastSinkSyncTime.set(0); + initStatsManager(); + } + + private ExecutionResult executeSingleServer() throws Exception { + Set entities = expandEntities(config.entities()); + batchSize.set(config.batchSize()); + + listeners.onJobConfigured(context, config); + + stats.set(initializeTotalRecords(entities)); + + String serverId = + org.openmetadata + .service + .apps + .bundles + .searchIndex + .distributed + .ServerIdentityResolver + .getInstance() + .getServerId(); + String jobId = + context.getJobId() != null ? context.getJobId().toString() : UUID.randomUUID().toString(); + this.failureRecorder = new IndexingFailureRecorder(collectionDAO, jobId, serverId); + cleanupOldFailures(); + + initializeSink(config); + + if (config.recreateIndex()) { + validateClusterCapacity(entities); + listeners.onIndexRecreationStarted(entities); + recreateContext = reCreateIndexes(entities); + } + + reIndexFromStartToEnd(entities); + closeSinkIfNeeded(); + // Promote anything yet to be promoted such as vector search indexes which is not part of + // entities set + finalizeReindex(); + + return buildResult(); + } + + private Set expandEntities(Set entities) { + if (entities.contains(ALL)) { + return getAll(); + } + return entities; + } + + private void validateClusterCapacity(Set entities) { + try { + SearchIndexClusterValidator validator = new SearchIndexClusterValidator(); + validator.validateCapacityForRecreate(searchRepository, entities); + } catch (InsufficientClusterCapacityException e) { + LOG.error("Cluster capacity check failed: {}", e.getMessage()); + throw e; + } catch (Exception e) { + LOG.warn("Failed to validate cluster capacity, proceeding with caution: {}", e.getMessage()); + } + } + + private void initializeSink(ReindexingConfiguration config) { + this.searchIndexSink = + searchRepository.createBulkSink( + config.batchSize(), config.maxConcurrentRequests(), config.payloadSize()); + this.recreateIndexHandler = searchRepository.createReindexHandler(); + + if (searchIndexSink != null) { + searchIndexSink.setFailureCallback(this::handleSinkFailure); + } + + LOG.debug("Initialized BulkSink with batch size: {}", config.batchSize()); + } + + private void handleSinkFailure( + String entityType, + String entityId, + String entityFqn, + String errorMessage, + IndexingFailureRecorder.FailureStage stage) { + if (failureRecorder != null) { + if (stage == IndexingFailureRecorder.FailureStage.PROCESS) { + failureRecorder.recordProcessFailure(entityType, entityId, entityFqn, errorMessage); + } else { + failureRecorder.recordSinkFailure(entityType, entityId, entityFqn, errorMessage); + } + } + } + + private void cleanupOldFailures() { + try { + long cutoffTime = System.currentTimeMillis() - TimeUnit.DAYS.toMillis(30); + int deleted = collectionDAO.searchIndexFailureDAO().deleteOlderThan(cutoffTime); + if (deleted > 0) { + LOG.info("Cleaned up {} old failure records", deleted); + } + } catch (Exception e) { + LOG.warn("Failed to cleanup old failure records", e); + } + } + + private void reIndexFromStartToEnd(Set entities) throws InterruptedException { + long totalEntities = + stats.get() != null && stats.get().getJobStats() != null + ? stats.get().getJobStats().getTotalRecords() + : 0; + + ThreadConfiguration threadConfig = calculateThreadConfiguration(totalEntities); + int effectiveQueueSize = initializeQueueAndExecutors(threadConfig, entities.size()); + + LOG.info( + "Starting reindexing with {} producers, {} consumers, queue size {}", + threadConfig.numProducers(), + threadConfig.numConsumers(), + effectiveQueueSize); + + executeReindexing(threadConfig.numConsumers(), entities); + } + + private ThreadConfiguration calculateThreadConfiguration(long totalEntities) { + int numConsumers = + config.consumerThreads() > 0 ? Math.min(config.consumerThreads(), MAX_CONSUMER_THREADS) : 2; + int numProducers = + config.producerThreads() > 1 + ? Math.min(config.producerThreads(), MAX_PRODUCER_THREADS) + : Math.clamp((int) (totalEntities / 10000), 2, MAX_PRODUCER_THREADS); + + return adjustThreadsForLimit(numProducers, numConsumers); + } + + private ThreadConfiguration adjustThreadsForLimit(int numProducers, int numConsumers) { + int entityCount = config.entities() != null ? config.entities().size() : 0; + int totalThreads = numProducers + numConsumers + entityCount; + + if (totalThreads > MAX_TOTAL_THREADS) { + LOG.warn( + "Total thread count {} exceeds limit {}, reducing...", totalThreads, MAX_TOTAL_THREADS); + double ratio = (double) MAX_TOTAL_THREADS / totalThreads; + numProducers = Math.max(1, (int) (numProducers * ratio)); + numConsumers = Math.max(1, (int) (numConsumers * ratio)); + } + + return new ThreadConfiguration(numProducers, numConsumers); + } + + private int initializeQueueAndExecutors(ThreadConfiguration threadConfig, int entityCount) { + int queueSize = config.queueSize() > 0 ? config.queueSize() : DEFAULT_QUEUE_SIZE; + int effectiveQueueSize = calculateMemoryAwareQueueSize(queueSize); + + taskQueue = new LinkedBlockingQueue<>(effectiveQueueSize); + producersDone.set(false); + + String jobIdTag = MDC.get("reindexJobId"); + String threadPrefix = "reindex-" + (jobIdTag != null ? jobIdTag + "-" : ""); + + int maxJobThreads = + Math.max(1, MAX_TOTAL_THREADS - threadConfig.numProducers() - threadConfig.numConsumers()); + int cappedEntityCount = Math.min(entityCount, maxJobThreads); + jobExecutor = + Executors.newFixedThreadPool( + cappedEntityCount, + Thread.ofPlatform() + .name(threadPrefix + "job-", 0) + .priority(Thread.MIN_PRIORITY) + .factory()); + + int finalNumConsumers = Math.min(threadConfig.numConsumers(), MAX_CONSUMER_THREADS); + consumerExecutor = + Executors.newFixedThreadPool( + finalNumConsumers, + Thread.ofPlatform() + .name(threadPrefix + "consumer-", 0) + .priority(Thread.MIN_PRIORITY) + .factory()); + + producerExecutor = + Executors.newFixedThreadPool( + threadConfig.numProducers(), + Thread.ofPlatform() + .name(threadPrefix + "producer-", 0) + .priority(Thread.MIN_PRIORITY) + .factory()); + + return effectiveQueueSize; + } + + private int calculateMemoryAwareQueueSize(int requestedSize) { + MemoryInfo memInfo = new MemoryInfo(); + long estimatedEntitySize = 5 * 1024L; + long maxQueueMemory = (long) (memInfo.maxMemory * 0.25); + long memoryBasedLimitLong = maxQueueMemory / (estimatedEntitySize * batchSize.get()); + int memoryBasedLimit = (int) Math.max(1, Math.min(memoryBasedLimitLong, Integer.MAX_VALUE)); + return Math.min(requestedSize, memoryBasedLimit); + } + + private void executeReindexing(int numConsumers, Set entities) + throws InterruptedException { + CountDownLatch consumerLatch = startConsumerThreads(numConsumers); + + try { + processEntityReindex(entities); + signalConsumersToStop(numConsumers); + waitForConsumersToComplete(consumerLatch); + } catch (InterruptedException e) { + LOG.info("Reindexing interrupted - stopping immediately"); + stopped.set(true); + Thread.currentThread().interrupt(); + throw e; + } finally { + cleanupExecutors(); + } + } + + private CountDownLatch startConsumerThreads(int numConsumers) { + CountDownLatch consumerLatch = new CountDownLatch(numConsumers); + Map mdc = MDC.getCopyOfContextMap(); + for (int i = 0; i < numConsumers; i++) { + final int consumerId = i; + consumerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + runConsumer(consumerId, consumerLatch); + } finally { + MDC.clear(); + } + }); + } + return consumerLatch; + } + + private void runConsumer(int consumerId, CountDownLatch consumerLatch) { + LOG.debug("Consumer {} started", consumerId); + try { + while (!stopped.get()) { + try { + IndexingTask task = taskQueue.poll(200, TimeUnit.MILLISECONDS); + if (task == null) { + continue; + } + if (POISON_PILL.equals(task.entityType())) { + break; + } + processTask(task); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + } finally { + LOG.debug("Consumer {} stopped", consumerId); + consumerLatch.countDown(); + } + } + + /** + * Process a single indexing task. + * + *

Stats are tracked via EntityStatsTracker (one per entity type) which flushes to + * search_index_server_stats table. Each stage tracks: + *

    + *
  • Reader: success/warnings/failed from ResultList + *
  • Process: success/failed during entity → search doc conversion (in BulkSink) + *
  • Sink: success/failed from ES/OS bulk response (in BulkSink) + *
  • Vector: success/failed for vector embeddings (in OpenSearchBulkSink) + *
+ */ + private void processTask(IndexingTask task) { + String entityType = task.entityType(); + ResultList entities = task.entities(); + Map contextData = createContextData(entityType); + EntityStatsTracker tracker = getTracker(entityType); + + // Stage 1: Reader stats (from source read) + int readerSuccessCount = listOrEmpty(entities.getData()).size(); + int readerFailedCount = listOrEmpty(entities.getErrors()).size(); + int readerWarningsCount = entities.getWarningsCount() != null ? entities.getWarningsCount() : 0; + + updateReaderStats(readerSuccessCount, readerFailedCount, readerWarningsCount); + if (tracker != null) { + tracker.recordReaderBatch(readerSuccessCount, readerFailedCount, readerWarningsCount); + } + + // Stage 2 & 3: Process + Sink handled by BulkSink via tracker passed in context + try { + writeEntitiesToSink(entityType, entities, contextData); + + StepStats currentEntityStats = createEntityStats(entities); + handleTaskSuccess(entityType, entities, currentEntityStats); + periodicSyncSinkStats(); + } catch (SearchIndexException e) { + handleSearchIndexException(entityType, entities, e); + } catch (Exception e) { + handleGenericException(entityType, entities, e); + } + } + + private Map createContextData(String entityType) { + return contextDataCache.computeIfAbsent( + entityType, + type -> { + Map contextData = new HashMap<>(); + contextData.put(ENTITY_TYPE_KEY, type); + contextData.put(RECREATE_INDEX, config.recreateIndex()); + contextData.put(RECREATE_CONTEXT, recreateContext); + contextData.put(BulkSink.STATS_TRACKER_CONTEXT_KEY, getSinkTracker(type)); + getTargetIndexForEntity(type) + .ifPresent(index -> contextData.put(TARGET_INDEX_KEY, index)); + return contextData; + }); + } + + private StageStatsTracker getSinkTracker(String entityType) { + if (context == null) { + return null; + } + return sinkTrackers.computeIfAbsent( + entityType, + et -> { + String jobId = context.getJobId().toString(); + String serverId = + org.openmetadata + .service + .apps + .bundles + .searchIndex + .distributed + .ServerIdentityResolver + .getInstance() + .getServerId(); + return new StageStatsTracker( + jobId, serverId, et, collectionDAO.searchIndexServerStatsDAO()); + }); + } + + private void writeEntitiesToSink( + String entityType, ResultList entities, Map contextData) throws Exception { + if (!TIME_SERIES_ENTITIES.contains(entityType)) { + @SuppressWarnings("unchecked") + List entityList = (List) entities.getData(); + searchIndexSink.write(entityList, contextData); + } else { + @SuppressWarnings("unchecked") + List entityList = + (List) entities.getData(); + searchIndexSink.write(entityList, contextData); + } + } + + private StepStats createEntityStats(ResultList entities) { + StepStats stepStats = new StepStats(); + stepStats.setSuccessRecords(listOrEmpty(entities.getData()).size()); + stepStats.setFailedRecords(listOrEmpty(entities.getErrors()).size()); + return stepStats; + } + + private void handleTaskSuccess( + String entityType, ResultList entities, StepStats currentEntityStats) { + if (entities.getErrors() != null && !entities.getErrors().isEmpty()) { + IndexingError error = + new IndexingError() + .withErrorSource(IndexingError.ErrorSource.READER) + .withSubmittedCount(batchSize.get()) + .withSuccessCount(entities.getData().size()) + .withFailedCount(entities.getErrors().size()) + .withMessage("Issues in Reading A Batch For Entities."); + listeners.onError(entityType, error, stats.get()); + } + + updateStats(entityType, currentEntityStats); + listeners.onProgressUpdate(stats.get(), context); + } + + private void handleSearchIndexException( + String entityType, ResultList entities, SearchIndexException e) { + if (!stopped.get()) { + IndexingError indexingError = e.getIndexingError(); + if (indexingError != null) { + listeners.onError(entityType, indexingError, stats.get()); + } else { + IndexingError error = createSinkError(e.getMessage()); + listeners.onError(entityType, error, stats.get()); + } + + syncSinkStatsFromBulkSink(); + + int dataSize = entities != null && entities.getData() != null ? entities.getData().size() : 0; + int readerErrors = entities != null ? listOrEmpty(entities.getErrors()).size() : 0; + StepStats failedStats = createFailedStats(indexingError, dataSize + readerErrors); + updateStats(entityType, failedStats); + } + LOG.error("Sink error for {}", entityType, e); + } + + private void handleGenericException(String entityType, ResultList entities, Exception e) { + if (!stopped.get()) { + IndexingError error = createSinkError(ExceptionUtils.getStackTrace(e)); + listeners.onError(entityType, error, stats.get()); + syncSinkStatsFromBulkSink(); + + int failedCount = + entities != null && entities.getData() != null ? entities.getData().size() : 0; + int readerErrors = entities != null ? listOrEmpty(entities.getErrors()).size() : 0; + StepStats failedStats = + new StepStats().withSuccessRecords(0).withFailedRecords(failedCount + readerErrors); + updateStats(entityType, failedStats); + } + LOG.error("Error for {}", entityType, e); + } + + private void signalConsumersToStop(int numConsumers) throws InterruptedException { + producersDone.set(true); + for (int i = 0; i < numConsumers; i++) { + taskQueue.put(new IndexingTask<>(POISON_PILL, null, -1)); + } + } + + private void waitForConsumersToComplete(CountDownLatch consumerLatch) + throws InterruptedException { + LOG.info("Waiting for consumers to complete..."); + consumerLatch.await(); + LOG.info("All consumers finished"); + } + + private void processEntityReindex(Set entities) throws InterruptedException { + // Use Phaser instead of pre-computed CountDownLatch to handle dynamic reader counts. + // Each entity type registers as a party, then dynamically registers its actual readers. + // This eliminates the batch-size-snapshot mismatch where auto-tune could desynchronize + // the pre-computed latch count from the actual number of readers created. + List ordered = EntityPriority.sortByPriority(entities); + LOG.info("Entity processing order: {}", ordered); + Phaser producerPhaser = new Phaser(entities.size()); + Map mdc = MDC.getCopyOfContextMap(); + + for (String entityType : ordered) { + jobExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + processEntityType(entityType, producerPhaser); + } finally { + MDC.clear(); + } + }); + } + + int phase = 0; + while (!producerPhaser.isTerminated()) { + if (stopped.get() || Thread.currentThread().isInterrupted()) { + LOG.info("Stop signal received during reindexing"); + if (producerExecutor != null) producerExecutor.shutdownNow(); + if (jobExecutor != null) jobExecutor.shutdownNow(); + return; + } + try { + producerPhaser.awaitAdvanceInterruptibly(phase, 1, TimeUnit.SECONDS); + break; + } catch (TimeoutException e) { + // Continue checking stop signal + } + } + } + + private void processEntityType(String entityType, Phaser producerPhaser) { + try { + int fixedBatchSize = EntityBatchSizeEstimator.estimateBatchSize(entityType, batchSize.get()); + int totalEntityRecords = getTotalEntityRecords(entityType); + listeners.onEntityTypeStarted(entityType, totalEntityRecords); + + entityBatchFailures.put(entityType, new AtomicInteger(0)); + + if (totalEntityRecords > 0) { + int numReaders = + Math.min( + calculateNumberOfThreads(totalEntityRecords, fixedBatchSize), + MAX_READERS_PER_ENTITY); + entityBatchCounters.put(entityType, new AtomicInteger(numReaders)); + + // Dynamically register actual readers with the phaser + producerPhaser.bulkRegister(numReaders); + + try { + if (TIME_SERIES_ENTITIES.contains(entityType)) { + Long filterStartTs = null; + Long filterEndTs = null; + if (config != null) { + long startTs = config.getTimeSeriesStartTs(entityType); + if (startTs > 0) { + filterStartTs = startTs; + filterEndTs = System.currentTimeMillis(); + } + } + final Long tsStart = filterStartTs; + final Long tsEnd = filterEndTs; + submitReaders( + entityType, + totalEntityRecords, + fixedBatchSize, + numReaders, + producerPhaser, + () -> { + PaginatedEntityTimeSeriesSource source = + (tsStart != null) + ? new PaginatedEntityTimeSeriesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalEntityRecords, + tsStart, + tsEnd) + : new PaginatedEntityTimeSeriesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalEntityRecords); + return source::readWithCursor; + }, + (readers, total) -> { + List cursors = new ArrayList<>(); + int perReader = total / readers; + for (int i = 1; i < readers; i++) { + cursors.add(RestUtil.encodeCursor(String.valueOf(i * perReader))); + } + return cursors; + }); + } else { + PaginatedEntitiesSource entSource = + new PaginatedEntitiesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalEntityRecords); + submitEntityReaders( + entityType, + totalEntityRecords, + fixedBatchSize, + numReaders, + producerPhaser, + entSource::findBoundaryCursors); + } + } catch (Exception e) { + LOG.error( + "Failed to submit readers for {}, deregistering {} phaser parties", + entityType, + numReaders, + e); + for (int i = 0; i < numReaders; i++) { + producerPhaser.arriveAndDeregister(); + } + throw e; + } + } else { + entityBatchCounters.put(entityType, new AtomicInteger(1)); + promoteEntityIndexIfReady(entityType); + } + + StepStats entityStats = + stats.get() != null && stats.get().getEntityStats() != null + ? stats.get().getEntityStats().getAdditionalProperties().get(entityType) + : null; + listeners.onEntityTypeCompleted(entityType, entityStats); + } catch (Exception e) { + LOG.error("Error processing entity type {}", entityType, e); + } finally { + // Deregister the entity coordinator party + producerPhaser.arriveAndDeregister(); + } + } + + private void submitReaders( + String entityType, + int totalRecords, + int fixedBatchSize, + int numReaders, + Phaser producerPhaser, + java.util.function.Supplier readerFactory, + java.util.function.BiFunction> boundaryFinder) { + Map mdc = MDC.getCopyOfContextMap(); + if (numReaders == 1) { + KeysetBatchReader reader = readerFactory.get(); + producerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + processKeysetBatches( + entityType, Integer.MAX_VALUE, fixedBatchSize, null, reader, producerPhaser); + } finally { + MDC.clear(); + } + }); + return; + } + + List boundaries = boundaryFinder.apply(numReaders, totalRecords); + int actualReaders = boundaries.size() + 1; + // Use ceiling division to avoid rounding-related entity loss at reader boundaries + int recordsPerReader = (totalRecords + actualReaders - 1) / actualReaders; + + if (actualReaders < numReaders) { + LOG.warn( + "Boundary discovery for {} returned {} cursors (expected {}), using {} readers", + entityType, + boundaries.size(), + numReaders - 1, + actualReaders); + entityBatchCounters.get(entityType).set(actualReaders); + // Deregister extra reader parties from the phaser + for (int j = 0; j < numReaders - actualReaders; j++) { + producerPhaser.arriveAndDeregister(); + } + } + + for (int i = 0; i < actualReaders; i++) { + String startCursor = (i == 0) ? null : boundaries.get(i - 1); + String endCursorForReader = (i < boundaries.size()) ? boundaries.get(i) : null; + int limit = (i == actualReaders - 1) ? Integer.MAX_VALUE : recordsPerReader; + KeysetBatchReader readerSource = readerFactory.get(); + final int readerLimit = limit; + final String readerEndCursor = endCursorForReader; + producerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + processKeysetBatches( + entityType, + readerLimit, + fixedBatchSize, + startCursor, + readerSource, + producerPhaser, + readerEndCursor); + } finally { + MDC.clear(); + } + }); + } + } + + @SuppressWarnings("unchecked") + private void submitEntityReaders( + String entityType, + int totalRecords, + int fixedBatchSize, + int numReaders, + Phaser producerPhaser, + java.util.function.BiFunction> boundaryFinder) { + Map mdc = MDC.getCopyOfContextMap(); + if (numReaders == 1) { + PaginatedEntitiesSource source = + new PaginatedEntitiesSource( + entityType, fixedBatchSize, getSearchIndexFields(entityType), totalRecords); + producerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + processKeysetBatches( + entityType, + Integer.MAX_VALUE, + fixedBatchSize, + null, + source::readNextKeyset, + producerPhaser); + } finally { + MDC.clear(); + } + }); + return; + } + + List boundaries = boundaryFinder.apply(numReaders, totalRecords); + int actualReaders = boundaries.size() + 1; + + if (actualReaders < numReaders) { + LOG.warn( + "Boundary discovery for {} returned {} cursors (expected {}), using {} readers", + entityType, + boundaries.size(), + numReaders - 1, + actualReaders); + entityBatchCounters.get(entityType).set(actualReaders); + for (int j = 0; j < numReaders - actualReaders; j++) { + producerPhaser.arriveAndDeregister(); + } + } + + for (int i = 0; i < actualReaders; i++) { + final String startCursor = (i == 0) ? null : boundaries.get(i - 1); + final boolean isLastReader = (i == actualReaders - 1); + + ListFilter filter; + if (isLastReader) { + filter = new ListFilter(Include.ALL); + } else { + String endBoundary = boundaries.get(i); + String decoded = RestUtil.decodeCursor(endBoundary); + Map cursorMap = + org.openmetadata.schema.utils.JsonUtils.readValue(decoded, Map.class); + filter = new BoundedListFilter(Include.ALL, cursorMap.get("name"), cursorMap.get("id")); + } + + final ListFilter readerFilter = filter; + producerExecutor.submit( + () -> { + if (mdc != null) MDC.setContextMap(mdc); + try { + PaginatedEntitiesSource source = + new PaginatedEntitiesSource( + entityType, + fixedBatchSize, + getSearchIndexFields(entityType), + totalRecords, + readerFilter); + processKeysetBatches( + entityType, + Integer.MAX_VALUE, + fixedBatchSize, + startCursor, + source::readNextKeyset, + producerPhaser); + } finally { + MDC.clear(); + } + }); + } + } + + private boolean hasReachedEndCursor(String afterCursor, String endCursor) { + if (endCursor == null || afterCursor == null) return false; + String decodedAfter = RestUtil.decodeCursor(afterCursor); + String decodedEnd = RestUtil.decodeCursor(endCursor); + if (decodedAfter == null || decodedEnd == null) return false; + + // Time-series cursors are numeric offsets + try { + int afterOffset = Integer.parseInt(decodedAfter); + int endOffset = Integer.parseInt(decodedEnd); + return afterOffset >= endOffset; + } catch (NumberFormatException ignored) { + // Not a numeric cursor, fall through to string comparison + } + return decodedAfter.equals(decodedEnd); + } + + private void processKeysetBatches( + String entityType, + int recordLimit, + int fixedBatchSize, + String startCursor, + KeysetBatchReader batchReader, + Phaser producerPhaser) { + processKeysetBatches( + entityType, recordLimit, fixedBatchSize, startCursor, batchReader, producerPhaser, null); + } + + private void processKeysetBatches( + String entityType, + int recordLimit, + int fixedBatchSize, + String startCursor, + KeysetBatchReader batchReader, + Phaser producerPhaser, + String endCursor) { + // Bypass the Redis-backed entity cache for the duration of this reader. Reindex never + // re-reads the same entity, so the cache hit rate is ~0; every relationship lookup pays a + // cache round-trip we don't need, and on an unhealthy Redis the indexer crawls because each + // miss pays a 300ms timeout. See {@link org.openmetadata.service.cache.EntityCacheBypass}. + try (org.openmetadata.service.cache.EntityCacheBypass.Handle ignored = + org.openmetadata.service.cache.EntityCacheBypass.skip()) { + processKeysetBatchesInternal( + entityType, + recordLimit, + fixedBatchSize, + startCursor, + batchReader, + producerPhaser, + endCursor); + } + } + + private void processKeysetBatchesInternal( + String entityType, + int recordLimit, + int fixedBatchSize, + String startCursor, + KeysetBatchReader batchReader, + Phaser producerPhaser, + String endCursor) { + boolean hadFailure = false; + try { + String keysetCursor = startCursor; + int processed = 0; + + while (processed < recordLimit && !stopped.get()) { + long backpressureWaitStart = System.currentTimeMillis(); + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 2000); + while (isBackpressureActive()) { + if (stopped.get()) { + return; + } + long elapsed = System.currentTimeMillis() - backpressureWaitStart; + if (elapsed > 15_000) { + LOG.warn("Backpressure wait timeout for {}, proceeding anyway", entityType); + break; + } + Thread.sleep(backoff.nextDelay()); + } + + try { + ResultList result = readWithRetry(batchReader, keysetCursor, entityType); + if (result == null || result.getData().isEmpty()) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (empty result)", + entityType, + processed, + recordLimit); + break; + } + + if (!stopped.get()) { + IndexingTask task = new IndexingTask<>(entityType, result, processed); + taskQueue.put(task); + } + + int readerSuccessCount = result.getData().size(); + int readerFailedCount = listOrEmpty(result.getErrors()).size(); + int readerWarningsCount = + result.getWarningsCount() != null ? result.getWarningsCount() : 0; + processed += readerSuccessCount + readerFailedCount + readerWarningsCount; + keysetCursor = result.getPaging() != null ? result.getPaging().getAfter() : null; + if (keysetCursor == null) { + LOG.debug( + "Reader for {} exhausted at processed={} of limit={} (null cursor)", + entityType, + processed, + recordLimit); + break; + } + if (hasReachedEndCursor(keysetCursor, endCursor)) { + LOG.debug("Reader for {} reached end cursor at processed={}", entityType, processed); + break; + } + } catch (SearchIndexException e) { + hadFailure = true; + LOG.error("Error reading keyset batch for {}", entityType, e); + if (failureRecorder != null) { + failureRecorder.recordReaderFailure( + entityType, e.getMessage(), ExceptionUtils.getStackTrace(e)); + } + listeners.onError(entityType, e.getIndexingError(), stats.get()); + int failedCount = + e.getIndexingError() != null && e.getIndexingError().getFailedCount() != null + ? e.getIndexingError().getFailedCount() + : fixedBatchSize; + updateReaderStats(0, failedCount, 0); + updateStats( + entityType, new StepStats().withSuccessRecords(0).withFailedRecords(failedCount)); + processed += fixedBatchSize; + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warn("Interrupted during keyset processing of {}", entityType); + } catch (Exception e) { + hadFailure = true; + if (!stopped.get()) { + LOG.error("Error in keyset processing for {}", entityType, e); + } + } finally { + producerPhaser.arriveAndDeregister(); + if (hadFailure) { + AtomicInteger failures = entityBatchFailures.get(entityType); + if (failures != null) { + failures.incrementAndGet(); + } + } + AtomicInteger remaining = entityBatchCounters.get(entityType); + if (remaining != null && remaining.decrementAndGet() == 0) { + promoteEntityIndexIfReady(entityType); + } + } + } + + private void processBatch(String entityType, int currentOffset, CountDownLatch producerLatch) { + // See note on processKeysetBatches: bypass the entity cache for reindex reader threads. + try (org.openmetadata.service.cache.EntityCacheBypass.Handle ignored = + org.openmetadata.service.cache.EntityCacheBypass.skip()) { + processBatchInternal(entityType, currentOffset, producerLatch); + } + } + + private void processBatchInternal( + String entityType, int currentOffset, CountDownLatch producerLatch) { + boolean batchHadFailure = false; + try { + if (stopped.get()) { + return; + } + + long backpressureWaitStart = System.currentTimeMillis(); + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 2000); + while (isBackpressureActive()) { + if (stopped.get()) { + return; + } + long elapsed = System.currentTimeMillis() - backpressureWaitStart; + if (elapsed > 15_000) { + LOG.warn( + "Backpressure wait timeout for {} offset {}, proceeding anyway", + entityType, + currentOffset); + break; + } + Thread.sleep(backoff.nextDelay()); + } + + Source source = createSource(entityType); + processReadTask(entityType, source, currentOffset); + } catch (Exception e) { + batchHadFailure = true; + if (!stopped.get()) { + LOG.error("Error processing batch for {}", entityType, e); + } + } finally { + producerLatch.countDown(); + // Track batch completion for per-entity promotion + if (batchHadFailure) { + AtomicInteger failures = entityBatchFailures.get(entityType); + if (failures != null) { + failures.incrementAndGet(); + } + } + AtomicInteger remaining = entityBatchCounters.get(entityType); + if (remaining != null && remaining.decrementAndGet() == 0) { + promoteEntityIndexIfReady(entityType); + } + } + } + + private void promoteEntityIndexIfReady(String entityType) { + if (recreateIndexHandler == null || recreateContext == null) { + return; + } + if (!config.recreateIndex()) { + return; + } + + if (!promotedEntities.add(entityType)) { + LOG.debug("Entity '{}' already promoted, skipping.", entityType); + return; + } + + AtomicInteger failures = entityBatchFailures.get(entityType); + boolean entitySuccess = failures == null || failures.get() == 0; + + Optional stagedIndexOpt = recreateContext.getStagedIndex(entityType); + if (stagedIndexOpt.isEmpty()) { + LOG.debug("No staged index found for entity '{}', skipping promotion.", entityType); + promotedEntities.remove(entityType); + return; + } + + EntityReindexContext entityContext = buildEntityReindexContext(entityType); + if (recreateIndexHandler instanceof DefaultRecreateHandler defaultHandler) { + LOG.info( + "Promoting index for entity '{}' (success={}, stagedIndex={})", + entityType, + entitySuccess, + stagedIndexOpt.get()); + defaultHandler.promoteEntityIndex(entityContext, entitySuccess); + + // When promoting the table index, also promote the column index since columns + // are indexed as part of table processing + if (Entity.TABLE.equals(entityType)) { + promoteColumnIndex(defaultHandler, entitySuccess); + } + } + } + + private void promoteColumnIndex(DefaultRecreateHandler handler, boolean tableSuccess) { + if (recreateContext == null) { + return; + } + Optional columnStagedIndex = recreateContext.getStagedIndex(Entity.TABLE_COLUMN); + if (columnStagedIndex.isEmpty()) { + return; + } + EntityReindexContext columnContext = buildEntityReindexContext(Entity.TABLE_COLUMN); + LOG.info( + "Promoting column index (success={}, stagedIndex={})", + tableSuccess, + columnStagedIndex.get()); + handler.promoteEntityIndex(columnContext, tableSuccess); + promotedEntities.add(Entity.TABLE_COLUMN); + } + + private ResultList readWithRetry( + KeysetBatchReader batchReader, String keysetCursor, String entityType) + throws SearchIndexException, InterruptedException { + int maxRetryAttempts = 3; + long retryBackoffMs = 500; + for (int attempt = 0; attempt <= maxRetryAttempts; attempt++) { + try { + return batchReader.readNextKeyset(keysetCursor); + } catch (SearchIndexException e) { + if (attempt >= maxRetryAttempts || !isTransientReadError(e)) { + throw e; + } + long backoffDelay = retryBackoffMs * (1L << attempt); + LOG.warn( + "Transient read failure for {} (attempt {}/{}), retrying in {}ms", + entityType, + attempt + 1, + maxRetryAttempts, + backoffDelay); + Thread.sleep(Math.min(backoffDelay, 10_000)); + } + } + return null; + } + + private boolean isTransientReadError(SearchIndexException e) { + String msg = e.getMessage(); + if (msg == null) { + msg = ""; + } + String lower = msg.toLowerCase(); + return lower.contains("timeout") + || lower.contains("connection") + || lower.contains("pool exhausted") + || lower.contains("connectexception") + || lower.contains("sockettimeoutexception") + || lower.contains("remotetransportexception"); + } + + private boolean isBackpressureActive() { + if (taskQueue != null) { + int size = taskQueue.size(); + int capacity = size + taskQueue.remainingCapacity(); + if (capacity > 0) { + int fillPercent = size * 100 / capacity; + ReindexingMetrics metrics = ReindexingMetrics.getInstance(); + if (metrics != null) { + metrics.updateQueueFillRatio(fillPercent); + } + return fillPercent > 90; + } + } + return false; + } + + private void processReadTask(String entityType, Source source, int offset) { + try { + if (stopped.get()) { + return; + } + + Object resultList = source.readWithCursor(RestUtil.encodeCursor(String.valueOf(offset))); + if (stopped.get()) { + return; + } + + if (resultList != null) { + ResultList entities = extractEntities(entityType, resultList); + if (!nullOrEmpty(entities.getData()) && !stopped.get()) { + IndexingTask task = new IndexingTask<>(entityType, entities, offset); + taskQueue.put(task); + } + } + } catch (SearchIndexException e) { + LOG.error("Error reading source for {}", entityType, e); + if (!stopped.get()) { + if (failureRecorder != null) { + failureRecorder.recordReaderFailure( + entityType, e.getMessage(), ExceptionUtils.getStackTrace(e)); + } + + listeners.onError(entityType, e.getIndexingError(), stats.get()); + IndexingError indexingError = e.getIndexingError(); + int failedCount = + indexingError != null && indexingError.getFailedCount() != null + ? indexingError.getFailedCount() + : batchSize.get(); + updateReaderStats(0, failedCount, 0); + StepStats failedStats = + new StepStats().withSuccessRecords(0).withFailedRecords(failedCount); + updateStats(entityType, failedStats); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOG.warn("Interrupted while queueing task for {}", entityType); + } + } + + private Source createSource(String entityType) { + String correctedEntityType = entityType; + if (QUERY_COST_RESULT_INCORRECT.equals(entityType)) { + LOG.warn(QUERY_COST_RESULT_WARNING); + correctedEntityType = QUERY_COST_RECORD; + } + + List searchIndexFields = getSearchIndexFields(correctedEntityType); + int knownTotal = getTotalEntityRecords(correctedEntityType); + + if (!TIME_SERIES_ENTITIES.contains(correctedEntityType)) { + return new PaginatedEntitiesSource( + correctedEntityType, batchSize.get(), searchIndexFields, knownTotal); + } else { + if (config != null) { + long startTs = config.getTimeSeriesStartTs(correctedEntityType); + if (startTs > 0) { + return new PaginatedEntityTimeSeriesSource( + correctedEntityType, + batchSize.get(), + searchIndexFields, + knownTotal, + startTs, + System.currentTimeMillis()); + } + } + return new PaginatedEntityTimeSeriesSource( + correctedEntityType, batchSize.get(), searchIndexFields, knownTotal); + } + } + + private List getSearchIndexFields(String entityType) { + // Delegate to the shared helper so single-server (this executor) and distributed + // (PartitionWorker) reindex paths request the same minimal field set. Otherwise + // setFieldsInBulk runs every registered fieldFetcher — including expensive ones like + // fetchAndSetOwns on Team/User — even though the search document drops most of them + // via getExcludedFields. PR #27723 originally fixed this for EntityReader; the + // executor was the missing piece. + return org.openmetadata.service.workflows.searchIndex.ReindexingUtil.getSearchIndexFields( + entityType); + } + + @SuppressWarnings("unchecked") + private ResultList extractEntities(String entityType, Object resultList) { + if (!TIME_SERIES_ENTITIES.contains(entityType)) { + return ((ResultList) resultList); + } else { + return ((ResultList) resultList); + } + } + + private Optional getTargetIndexForEntity(String entityType) { + if (recreateContext == null) { + return Optional.empty(); + } + + Optional stagedIndex = recreateContext.getStagedIndex(entityType); + if (stagedIndex.isPresent()) { + return stagedIndex; + } + + if (QUERY_COST_RESULT_INCORRECT.equals(entityType)) { + return recreateContext.getStagedIndex(QUERY_COST_RECORD); + } + + return Optional.empty(); + } + + public Stats initializeTotalRecords(Set entities) { + Stats jobDataStats = new Stats(); + jobDataStats.setEntityStats(new EntityStats()); + + int total = 0; + for (String entityType : entities) { + int entityTotal = getEntityTotal(entityType); + total += entityTotal; + + StepStats entityStats = new StepStats(); + entityStats.setTotalRecords(entityTotal); + entityStats.setSuccessRecords(0); + entityStats.setFailedRecords(0); + + jobDataStats.getEntityStats().getAdditionalProperties().put(entityType, entityStats); + } + + StepStats jobStats = new StepStats(); + jobStats.setTotalRecords(total); + jobStats.setSuccessRecords(0); + jobStats.setFailedRecords(0); + jobDataStats.setJobStats(jobStats); + + StepStats readerStats = new StepStats(); + readerStats.setTotalRecords(total); + readerStats.setSuccessRecords(0); + readerStats.setFailedRecords(0); + readerStats.setWarningRecords(0); + jobDataStats.setReaderStats(readerStats); + + StepStats sinkStats = new StepStats(); + sinkStats.setTotalRecords(0); + sinkStats.setSuccessRecords(0); + sinkStats.setFailedRecords(0); + jobDataStats.setSinkStats(sinkStats); + + StepStats processStats = new StepStats(); + processStats.setTotalRecords(0); + processStats.setSuccessRecords(0); + processStats.setFailedRecords(0); + jobDataStats.setProcessStats(processStats); + + // Add a stats slot for TABLE_COLUMN since columns are indexed as part of table processing + // but TABLE_COLUMN is not a standalone entity in the entities set + if (entities.contains(Entity.TABLE) && !entities.contains(Entity.TABLE_COLUMN)) { + StepStats columnEntityStats = new StepStats(); + columnEntityStats.setTotalRecords(0); + columnEntityStats.setSuccessRecords(0); + columnEntityStats.setFailedRecords(0); + jobDataStats + .getEntityStats() + .getAdditionalProperties() + .put(Entity.TABLE_COLUMN, columnEntityStats); + LOG.info("Added TABLE_COLUMN stats slot for column indexing tracking"); + } + + return jobDataStats; + } + + private int getEntityTotal(String entityType) { + try { + String correctedEntityType = entityType; + if (QUERY_COST_RESULT_INCORRECT.equals(entityType)) { + LOG.warn(QUERY_COST_RESULT_WARNING); + correctedEntityType = QUERY_COST_RECORD; + } + + if (!TIME_SERIES_ENTITIES.contains(correctedEntityType)) { + EntityRepository repository = Entity.getEntityRepository(correctedEntityType); + return repository.getDao().listCount(new ListFilter(Include.ALL)); + } else { + EntityTimeSeriesRepository repository; + ListFilter listFilter = new ListFilter(null); + if (isDataInsightIndex(entityType)) { + listFilter.addQueryParam("entityFQNHash", FullyQualifiedName.buildHash(entityType)); + repository = Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA); + } else { + repository = Entity.getEntityTimeSeriesRepository(entityType); + } + if (config != null) { + long startTs = config.getTimeSeriesStartTs(correctedEntityType); + if (startTs > 0) { + long endTs = System.currentTimeMillis(); + return repository.getTimeSeriesDao().listCount(listFilter, startTs, endTs, false); + } + } + return repository.getTimeSeriesDao().listCount(listFilter); + } + } catch (Exception e) { + LOG.debug("Error getting total for '{}'", entityType, e); + return 0; + } + } + + private int getTotalEntityRecords(String entityType) { + if (stats.get() == null + || stats.get().getEntityStats() == null + || stats.get().getEntityStats().getAdditionalProperties() == null) { + return 0; + } + + StepStats entityStats = stats.get().getEntityStats().getAdditionalProperties().get(entityType); + if (entityStats != null) { + return entityStats.getTotalRecords() != null ? entityStats.getTotalRecords() : 0; + } + return 0; + } + + private int calculateNumberOfThreads(int totalEntityRecords, int fixedBatchSize) { + if (fixedBatchSize <= 0) return 1; + int mod = totalEntityRecords % fixedBatchSize; + if (mod == 0) { + return totalEntityRecords / fixedBatchSize; + } else { + return (totalEntityRecords / fixedBatchSize) + 1; + } + } + + // Stats is published once via stats.set(initializeTotalRecords(...)) and all subsequent + // mutations operate on that same mutable object under synchronized methods. + + synchronized void updateStats(String entityType, StepStats currentEntityStats) { + Stats jobDataStats = stats.get(); + if (jobDataStats == null) { + return; + } + + updateEntityStats(jobDataStats, entityType, currentEntityStats); + + // When processing tables, also update column stats from the sink + if (Entity.TABLE.equals(entityType) && searchIndexSink != null) { + updateColumnStatsFromSink(jobDataStats); + } + + updateJobStats(jobDataStats); + } + + private void updateColumnStatsFromSink(Stats jobDataStats) { + if (searchIndexSink == null || jobDataStats == null || jobDataStats.getEntityStats() == null) { + return; + } + StepStats columnStats = searchIndexSink.getColumnStats(); + if (columnStats != null && columnStats.getTotalRecords() > 0) { + StepStats existingColumnStats = + jobDataStats.getEntityStats().getAdditionalProperties().get(Entity.TABLE_COLUMN); + if (existingColumnStats != null) { + existingColumnStats.setTotalRecords(columnStats.getTotalRecords()); + existingColumnStats.setSuccessRecords(columnStats.getSuccessRecords()); + existingColumnStats.setFailedRecords(columnStats.getFailedRecords()); + } + } + } + + synchronized void updateReaderStats(int successCount, int failedCount, int warningsCount) { + Stats jobDataStats = stats.get(); + if (jobDataStats == null) { + return; + } + + StepStats readerStats = jobDataStats.getReaderStats(); + if (readerStats == null) { + readerStats = new StepStats(); + jobDataStats.setReaderStats(readerStats); + } + + int currentSuccess = + readerStats.getSuccessRecords() != null ? readerStats.getSuccessRecords() : 0; + int currentFailed = readerStats.getFailedRecords() != null ? readerStats.getFailedRecords() : 0; + int currentWarnings = + readerStats.getWarningRecords() != null ? readerStats.getWarningRecords() : 0; + + readerStats.setSuccessRecords(currentSuccess + successCount); + readerStats.setFailedRecords(currentFailed + failedCount); + readerStats.setWarningRecords(currentWarnings + warningsCount); + } + + synchronized void updateSinkTotalSubmitted(int submittedCount) { + Stats jobDataStats = stats.get(); + if (jobDataStats == null) { + return; + } + + StepStats sinkStats = jobDataStats.getSinkStats(); + if (sinkStats == null) { + sinkStats = new StepStats(); + sinkStats.setTotalRecords(0); + jobDataStats.setSinkStats(sinkStats); + } + + int currentTotal = sinkStats.getTotalRecords() != null ? sinkStats.getTotalRecords() : 0; + sinkStats.setTotalRecords(currentTotal + submittedCount); + } + + synchronized void syncSinkStatsFromBulkSink() { + if (searchIndexSink == null) { + return; + } + + Stats jobDataStats = stats.get(); + if (jobDataStats == null) { + return; + } + + StepStats bulkSinkStats = searchIndexSink.getStats(); + if (bulkSinkStats == null) { + return; + } + + StepStats sinkStats = jobDataStats.getSinkStats(); + if (sinkStats == null) { + sinkStats = new StepStats(); + jobDataStats.setSinkStats(sinkStats); + } + + sinkStats.setTotalRecords( + bulkSinkStats.getTotalRecords() != null ? bulkSinkStats.getTotalRecords() : 0); + sinkStats.setSuccessRecords( + bulkSinkStats.getSuccessRecords() != null ? bulkSinkStats.getSuccessRecords() : 0); + sinkStats.setFailedRecords( + bulkSinkStats.getFailedRecords() != null ? bulkSinkStats.getFailedRecords() : 0); + + // Sync vector stats if available + StepStats vectorStats = searchIndexSink.getVectorStats(); + if (vectorStats != null + && (vectorStats.getTotalRecords() != null && vectorStats.getTotalRecords() > 0)) { + jobDataStats.setVectorStats(vectorStats); + } + + // Sync process stats if available + StepStats processStats = searchIndexSink.getProcessStats(); + if (processStats != null) { + jobDataStats.setProcessStats(processStats); + } + } + + private void periodicSyncSinkStats() { + long now = System.currentTimeMillis(); + long last = lastSinkSyncTime.get(); + if (now - last >= SINK_SYNC_INTERVAL_MS && lastSinkSyncTime.compareAndSet(last, now)) { + syncSinkStatsFromBulkSink(); + } + } + + private void updateEntityStats(Stats statsObj, String entityType, StepStats currentEntityStats) { + if (statsObj.getEntityStats() == null + || statsObj.getEntityStats().getAdditionalProperties() == null) { + return; + } + + StepStats entityStats = statsObj.getEntityStats().getAdditionalProperties().get(entityType); + if (entityStats != null) { + entityStats.withSuccessRecords( + entityStats.getSuccessRecords() + currentEntityStats.getSuccessRecords()); + entityStats.withFailedRecords( + entityStats.getFailedRecords() + currentEntityStats.getFailedRecords()); + + int actual = entityStats.getSuccessRecords() + entityStats.getFailedRecords(); + if (actual > entityStats.getTotalRecords()) { + entityStats.setTotalRecords(actual); + } + } + } + + private void updateJobStats(Stats statsObj) { + StepStats jobStats = statsObj.getJobStats(); + if (jobStats == null || statsObj.getEntityStats() == null) { + return; + } + + int totalRecords = + statsObj.getEntityStats().getAdditionalProperties().entrySet().stream() + .filter(e -> !Entity.TABLE_COLUMN.equals(e.getKey())) + .mapToInt(e -> e.getValue().getTotalRecords()) + .sum(); + + int totalSuccess = + statsObj.getEntityStats().getAdditionalProperties().entrySet().stream() + .filter(e -> !Entity.TABLE_COLUMN.equals(e.getKey())) + .mapToInt(e -> e.getValue().getSuccessRecords()) + .sum(); + + int totalFailed = + statsObj.getEntityStats().getAdditionalProperties().entrySet().stream() + .filter(e -> !Entity.TABLE_COLUMN.equals(e.getKey())) + .mapToInt(e -> e.getValue().getFailedRecords()) + .sum(); + + jobStats + .withTotalRecords(totalRecords) + .withSuccessRecords(totalSuccess) + .withFailedRecords(totalFailed); + + StepStats readerStats = statsObj.getReaderStats(); + if (readerStats != null && totalRecords > readerStats.getTotalRecords()) { + readerStats.setTotalRecords(totalRecords); + } + } + + private IndexingError createSinkError(String message) { + return new IndexingError().withErrorSource(IndexingError.ErrorSource.SINK).withMessage(message); + } + + private StepStats createFailedStats(IndexingError indexingError, int dataSize) { + StepStats failedStats = new StepStats(); + failedStats.setSuccessRecords(indexingError != null ? indexingError.getSuccessCount() : 0); + failedStats.setFailedRecords(indexingError != null ? indexingError.getFailedCount() : dataSize); + return failedStats; + } + + private Set getAll() { + return new HashSet<>(searchRepository.getEntityIndexMap().keySet()); + } + + private ReindexContext reCreateIndexes(Set entities) { + if (recreateIndexHandler == null) { + return null; + } + return recreateIndexHandler.reCreateIndexes(entities); + } + + private void closeSinkIfNeeded() { + if (searchIndexSink != null && sinkClosed.compareAndSet(false, true)) { + int pendingVectorTasks = searchIndexSink.getPendingVectorTaskCount(); + if (pendingVectorTasks > 0) { + LOG.info( + "Waiting for {} pending vector embedding tasks to complete before closing", + pendingVectorTasks); + VectorCompletionResult vcResult = searchIndexSink.awaitVectorCompletionWithDetails(300); + LOG.info( + "Vector completion: completed={}, pending={}, waited={}ms", + vcResult.completed(), + vcResult.pendingTaskCount(), + vcResult.waitedMillis()); + } + + LOG.info("Forcing final flush of bulk processor and vector embeddings"); + searchIndexSink.close(); + syncSinkStatsFromBulkSink(); + } + } + + private ExecutionResult buildResult() { + if (failureRecorder != null) { + failureRecorder.flush(); + } + + syncSinkStatsFromBulkSink(); + updateColumnStatsFromSink(stats.get()); + + Stats currentStats = stats.get(); + if (currentStats != null) { + StatsReconciler.reconcile(currentStats); + } + + long endTime = System.currentTimeMillis(); + ExecutionResult.Status status = determineStatus(); + + if (status == ExecutionResult.Status.COMPLETED) { + listeners.onJobCompleted(stats.get(), endTime - startTime); + } else if (status == ExecutionResult.Status.COMPLETED_WITH_ERRORS) { + listeners.onJobCompletedWithErrors(stats.get(), endTime - startTime); + } else if (status == ExecutionResult.Status.STOPPED) { + listeners.onJobStopped(stats.get()); + } + + return ExecutionResult.fromStats(stats.get(), status, startTime); + } + + private ExecutionResult.Status determineStatus() { + if (stopped.get()) { + return ExecutionResult.Status.STOPPED; + } + + if (hasIncompleteProcessing()) { + return ExecutionResult.Status.COMPLETED_WITH_ERRORS; + } + + return ExecutionResult.Status.COMPLETED; + } + + private boolean hasIncompleteProcessing() { + Stats currentStats = stats.get(); + if (currentStats == null || currentStats.getJobStats() == null) { + return false; + } + + StepStats jobStats = currentStats.getJobStats(); + long failed = jobStats.getFailedRecords() != null ? jobStats.getFailedRecords() : 0; + long processed = jobStats.getSuccessRecords() != null ? jobStats.getSuccessRecords() : 0; + long total = jobStats.getTotalRecords() != null ? jobStats.getTotalRecords() : 0; + + return failed > 0 || (total > 0 && processed < total); + } + + public void stop() { + LOG.info("Stopping reindexing executor..."); + stopped.set(true); + producersDone.set(true); + + listeners.onJobStopped(stats.get()); + + if (searchIndexSink != null) { + LOG.info( + "Stopping executor: flushing sink ({} active bulk requests)", + searchIndexSink.getActiveBulkRequestCount()); + searchIndexSink.flushAndAwait(10); + } + + int dropped = taskQueue != null ? taskQueue.size() : 0; + if (dropped > 0) { + LOG.warn("Dropping {} queued tasks during shutdown", dropped); + } + + shutdownExecutor(producerExecutor, "producer"); + shutdownExecutor(jobExecutor, "job"); + + if (taskQueue != null) { + taskQueue.clear(); + for (int i = 0; i < MAX_CONSUMER_THREADS; i++) { + taskQueue.offer(new IndexingTask<>(POISON_PILL, null, -1)); + } + } + if (consumerExecutor != null && !consumerExecutor.isShutdown()) { + consumerExecutor.shutdown(); + try { + if (!consumerExecutor.awaitTermination(5, TimeUnit.SECONDS)) { + consumerExecutor.shutdownNow(); + LOG.warn("Consumer executor did not terminate within 5s, forced shutdown"); + } + } catch (InterruptedException e) { + consumerExecutor.shutdownNow(); + Thread.currentThread().interrupt(); + } + } + + LOG.info("Reindexing executor stopped"); + } + + public boolean isStopped() { + return stopped.get(); + } + + private void cleanupExecutors() { + if (!stopped.get()) { + shutdownExecutor(consumerExecutor, "consumer", 30, TimeUnit.SECONDS); + shutdownExecutor(jobExecutor, "job", 20, TimeUnit.SECONDS); + shutdownExecutor(producerExecutor, "producer", 1, TimeUnit.MINUTES); + } + } + + private void shutdownExecutor(ExecutorService executor, String name) { + if (executor != null && !executor.isShutdown()) { + LOG.info("Force shutting down {} executor", name); + List pendingTasks = executor.shutdownNow(); + LOG.info("Cancelled {} pending {} tasks", pendingTasks.size(), name); + } + } + + private void shutdownExecutor( + ExecutorService executor, String name, long timeout, TimeUnit unit) { + if (executor != null && !executor.isShutdown()) { + executor.shutdown(); + try { + if (!executor.awaitTermination(timeout, unit)) { + executor.shutdownNow(); + LOG.warn("{} did not terminate within timeout", name); + } + } catch (InterruptedException e) { + executor.shutdownNow(); + Thread.currentThread().interrupt(); + } + } + } + + private void cleanup() { + if (failureRecorder != null) { + try { + failureRecorder.close(); + } catch (Exception e) { + LOG.error("Error closing failure recorder", e); + } + } + + if (searchIndexSink != null && sinkClosed.compareAndSet(false, true)) { + try { + searchIndexSink.close(); + } catch (Exception e) { + LOG.error("Error closing search index sink", e); + } + } + + finalizeReindex(); + } + + private void finalizeReindex() { + if (recreateIndexHandler == null || recreateContext == null) { + return; + } + + try { + recreateContext + .getEntities() + .forEach( + entityType -> { + // Skip entities already promoted via per-entity promotion + if (promotedEntities.contains(entityType)) { + LOG.debug( + "Skipping finalizeReindex for entity '{}' - already promoted.", entityType); + return; + } + try { + AtomicInteger failures = entityBatchFailures.get(entityType); + boolean entitySuccess = + !stopped.get() && (failures == null || failures.get() == 0); + recreateIndexHandler.finalizeReindex( + buildEntityReindexContext(entityType), entitySuccess); + } catch (Exception ex) { + LOG.error("Failed to finalize reindex for {}", entityType, ex); + } + }); + } finally { + recreateContext = null; + promotedEntities.clear(); + } + } + + private EntityReindexContext buildEntityReindexContext(String entityType) { + return EntityReindexContext.builder() + .entityType(entityType) + .originalIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .canonicalIndex(recreateContext.getCanonicalIndex(entityType).orElse(null)) + .activeIndex(recreateContext.getOriginalIndex(entityType).orElse(null)) + .stagedIndex(recreateContext.getStagedIndex(entityType).orElse(null)) + .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) + .existingAliases(recreateContext.getExistingAliases(entityType)) + .parentAliases(new HashSet<>(listOrEmpty(recreateContext.getParentAliases(entityType)))) + .build(); + } + + @Override + public void close() { + if (statsManager != null) { + statsManager.flushAll(); + } + sinkTrackers.values().forEach(StageStatsTracker::flush); + stop(); + cleanup(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategy.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategy.java new file mode 100644 index 000000000000..d347514bdb69 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategy.java @@ -0,0 +1,41 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import java.util.Optional; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.search.SearchRepository; + +public class SingleServerIndexingStrategy implements IndexingStrategy { + + private final SearchIndexExecutor executor; + + public SingleServerIndexingStrategy( + CollectionDAO collectionDAO, SearchRepository searchRepository) { + this.executor = new SearchIndexExecutor(collectionDAO, searchRepository); + } + + @Override + public void addListener(ReindexingProgressListener listener) { + executor.addListener(listener); + } + + @Override + public ExecutionResult execute(ReindexingConfiguration config, ReindexingJobContext context) { + return executor.execute(config, context); + } + + @Override + public Optional getStats() { + return Optional.ofNullable(executor.getStats().get()); + } + + @Override + public void stop() { + executor.stop(); + } + + @Override + public boolean isStopped() { + return executor.isStopped(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DISTRIBUTED_INDEXING.md b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DISTRIBUTED_INDEXING.md index 354fb9ed2091..3d22f758a985 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DISTRIBUTED_INDEXING.md +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DISTRIBUTED_INDEXING.md @@ -250,23 +250,23 @@ WHERE lockKey = ? AND jobId = ? ## Configuration -Distributed indexing is always enabled. Tune the reindex API like so: +Enable distributed indexing via the reindex API: ```json { "entities": ["table", "database", "topic", "dashboard"], + "recreateIndex": true, "batchSize": 100, - "consumerThreads": 4 + "consumerThreads": 4, + "useDistributedIndexing": true } ``` -Search indexing always writes to staged indexes and promotes aliases after successful processing so -live search indexes are not mutated during the bulk rebuild. - ### Configuration Options | Parameter | Default | Description | |-----------|---------|-------------| +| useDistributedIndexing | false | Enable distributed mode | | batchSize | 100 | Entities per batch | | consumerThreads | 4 | Worker threads per server | | maxConcurrentRequests | 100 | Concurrent ES/OS requests | diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContext.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContext.java index 6e05f162daec..fd6677781744 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContext.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContext.java @@ -57,6 +57,11 @@ public UUID getAppId() { return job.getId(); } + @Override + public boolean isDistributed() { + return true; + } + @Override public String getSource() { return source; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifier.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifier.java index 1f8b2f4b49e2..7b0e26dd6da5 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifier.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifier.java @@ -18,6 +18,13 @@ /** * Interface for notifying servers about distributed job events. + * + *

This abstraction allows for different notification mechanisms: + * + *

    + *
  • Redis Pub/Sub - instant push notifications when Redis is available + *
  • Database polling - fallback when Redis is not configured + *
*/ public interface DistributedJobNotifier { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactory.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactory.java index df288fe0670c..b90685d0f8a2 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactory.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactory.java @@ -14,10 +14,14 @@ package org.openmetadata.service.apps.bundles.searchIndex.distributed; import lombok.extern.slf4j.Slf4j; +import org.openmetadata.service.cache.CacheConfig; import org.openmetadata.service.jdbi3.CollectionDAO; /** - * Factory for creating the DistributedJobNotifier used by search indexing. + * Factory for creating the appropriate DistributedJobNotifier based on configuration. + * + *

Uses Redis Pub/Sub when Redis is configured and available, otherwise falls back to database + * polling. */ @Slf4j public class DistributedJobNotifierFactory { @@ -27,14 +31,42 @@ private DistributedJobNotifierFactory() { } /** - * Create a DistributedJobNotifier. + * Create a DistributedJobNotifier based on the current configuration. * + * @param cacheConfig The cache configuration (contains Redis settings) * @param collectionDAO The DAO for database access * @param serverId The current server's ID - * @return The notifier implementation + * @return The appropriate notifier implementation */ - public static DistributedJobNotifier create(CollectionDAO collectionDAO, String serverId) { - LOG.info("Using database polling for distributed search indexing job discovery"); + public static DistributedJobNotifier create( + CacheConfig cacheConfig, CollectionDAO collectionDAO, String serverId) { + + if (cacheConfig != null && cacheConfig.provider == CacheConfig.Provider.redis) { + // Redis is configured - try to use Redis Pub/Sub + if (isRedisConfigValid(cacheConfig)) { + LOG.info( + "Redis is configured - using Redis Pub/Sub for distributed job notifications (instant discovery)"); + return new RedisJobNotifier(cacheConfig, serverId); + } else { + LOG.warn( + "Redis is configured but URL is missing - falling back to database polling for job notifications"); + } + } + + LOG.info( + "Redis not configured - using database polling for distributed job notifications (30s discovery delay)"); return new PollingJobNotifier(collectionDAO, serverId); } + + /** + * Check if Redis configuration is valid and complete. + * + * @param cacheConfig The cache configuration + * @return true if Redis can be used + */ + private static boolean isRedisConfigValid(CacheConfig cacheConfig) { + return cacheConfig.redis != null + && cacheConfig.redis.url != null + && !cacheConfig.redis.url.isEmpty(); + } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java index f2c407b0ce54..c0dc8ce20b43 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java @@ -25,9 +25,9 @@ import org.openmetadata.service.Entity; import org.openmetadata.service.apps.bundles.searchIndex.BulkSink; import org.openmetadata.service.apps.bundles.searchIndex.IndexingFailureRecorder; +import org.openmetadata.service.cache.CacheConfig; import org.openmetadata.service.jdbi3.AppRepository; import org.openmetadata.service.jdbi3.CollectionDAO; -import org.openmetadata.service.search.ReindexContext; import org.openmetadata.service.search.SearchClusterMetrics; import org.openmetadata.service.search.SearchRepository; @@ -39,7 +39,12 @@ * service runs on all servers and allows non-triggering servers to discover and participate in * active jobs. * - *

Job discovery is handled by a {@link DistributedJobNotifier} backed by database polling. + *

Job discovery is handled by a {@link DistributedJobNotifier}: + * + *

    + *
  • When Redis is configured: Uses Redis Pub/Sub for instant notification + *
  • When Redis is not available: Falls back to database polling (30s interval) + *
*/ @Slf4j public class DistributedJobParticipant implements Managed { @@ -68,12 +73,15 @@ public class DistributedJobParticipant implements Managed { private volatile Thread participantThread; public DistributedJobParticipant( - CollectionDAO collectionDAO, SearchRepository searchRepository, String serverId) { + CollectionDAO collectionDAO, + SearchRepository searchRepository, + String serverId, + CacheConfig cacheConfig) { this( collectionDAO, searchRepository, serverId, - DistributedJobNotifierFactory.create(collectionDAO, serverId)); + DistributedJobNotifierFactory.create(cacheConfig, collectionDAO, serverId)); } /** @@ -103,7 +111,7 @@ public void start() { // Register callback to receive job start notifications notifier.onJobStarted(this::onJobDiscovered); - // Start the notifier + // Start the notifier (Redis subscription or polling) notifier.start(); // Start orphan job monitor to detect jobs left behind by crashed coordinators @@ -181,16 +189,7 @@ private void onJobDiscovered(UUID jobId) { // Check if there are pending partitions we can help with long pendingCount = coordinator.getPartitions(job.getId(), PartitionStatus.PENDING).size(); if (pendingCount == 0) { - long processingCount = - coordinator.getPartitions(job.getId(), PartitionStatus.PROCESSING).size(); - long completedCount = - coordinator.getPartitions(job.getId(), PartitionStatus.COMPLETED).size(); - LOG.info( - "Discovered distributed job {} on server {}, but no pending partitions remain (processing={}, completed={}); not joining", - job.getId(), - serverId, - processingCount, - completedCount); + LOG.debug("No pending partitions to process for job {}", job.getId()); return; } @@ -306,12 +305,6 @@ private void processJobPartitions(SearchIndexJob job) { DistributedJobStatsAggregator statsAggregator = null; AppRunRecordContext appCtx = null; try { - Optional stagedIndexContext = buildStagedIndexContext(job); - if (stagedIndexContext.isEmpty()) { - return; - } - ReindexContext reindexContext = stagedIndexContext.orElseThrow(); - appCtx = resolveAppRunRecordContext(); if (appCtx != null) { restoreAppRunRecordToRunning(appCtx.appId(), appCtx.startTime()); @@ -348,6 +341,22 @@ private void processJobPartitions(SearchIndexJob job) { ? job.getJobConfiguration().getBatchSize() : 100; + // Check if this job is doing index recreation + boolean recreateIndex = Boolean.TRUE.equals(job.getJobConfiguration().getRecreateIndex()); + org.openmetadata.service.search.ReindexContext recreateContext = null; + + if (recreateIndex && job.getStagedIndexMapping() != null) { + // Reconstruct context from job's staged index mapping + recreateContext = + org.openmetadata.service.search.ReindexContext.fromStagedIndexMapping( + job.getStagedIndexMapping()); + LOG.info( + "Participant using staged index mapping from job {}: {}", + job.getId(), + job.getStagedIndexMapping()); + } + + // Set up failure callback on bulk sink to record sink failures final IndexingFailureRecorder recorder = failureRecorder; bulkSink.setFailureCallback( (entityType, entityId, entityFqn, errorMessage, stage) -> { @@ -360,8 +369,10 @@ private void processJobPartitions(SearchIndexJob job) { } }); + // Create partition worker with recreate context and failure recorder PartitionWorker worker = - new PartitionWorker(coordinator, bulkSink, batchSize, reindexContext, failureRecorder); + new PartitionWorker( + coordinator, bulkSink, batchSize, recreateContext, recreateIndex, failureRecorder); int partitionsProcessed = 0; long totalReaderSuccess = 0; @@ -475,21 +486,6 @@ private void processJobPartitions(SearchIndexJob job) { } } - private Optional buildStagedIndexContext(SearchIndexJob job) { - if (job.getStagedIndexMapping() == null || job.getStagedIndexMapping().isEmpty()) { - LOG.warn( - "Skipping distributed reindex job {} on server {} because staged index mapping is missing", - job.getId(), - serverId); - return Optional.empty(); - } - LOG.info( - "Participant using staged index mapping from job {}: {}", - job.getId(), - job.getStagedIndexMapping()); - return Optional.of(ReindexContext.fromStagedIndexMapping(job.getStagedIndexMapping())); - } - /** Check if currently participating in a job. */ public boolean isParticipating() { return participating.get(); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java index e0e22263b0ef..91b30bb593b1 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexCoordinator.java @@ -29,7 +29,6 @@ import org.openmetadata.schema.utils.JsonUtils; import org.openmetadata.service.Entity; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; -import org.openmetadata.service.apps.bundles.searchIndex.SearchIndexEntityTypes; import org.openmetadata.service.jdbi3.CollectionDAO; import org.openmetadata.service.jdbi3.CollectionDAO.SearchIndexJobDAO; import org.openmetadata.service.jdbi3.CollectionDAO.SearchIndexJobDAO.SearchIndexJobRecord; @@ -246,7 +245,7 @@ private void precomputePartitionStartCursors(UUID jobId, List> byEntity = partitions.stream() .filter(p -> p.getEntityType() != null) - .filter(p -> !SearchIndexEntityTypes.isTimeSeriesEntity(p.getEntityType())) + .filter(p -> !PartitionWorker.TIME_SERIES_ENTITIES.contains(p.getEntityType())) .collect(Collectors.groupingBy(SearchIndexPartition::getEntityType)); Map> jobCache = new HashMap<>(); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java index efeb6ea84f08..0950a0c490eb 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java @@ -13,9 +13,12 @@ package org.openmetadata.service.apps.bundles.searchIndex.distributed; +import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; + import io.micrometer.core.instrument.Timer; import java.util.ArrayList; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -36,7 +39,6 @@ import org.openmetadata.service.apps.bundles.searchIndex.BulkSink; import org.openmetadata.service.apps.bundles.searchIndex.CompositeProgressListener; import org.openmetadata.service.apps.bundles.searchIndex.ElasticSearchBulkSink; -import org.openmetadata.service.apps.bundles.searchIndex.EntityReindexContextMapper; import org.openmetadata.service.apps.bundles.searchIndex.IndexingFailureRecorder; import org.openmetadata.service.apps.bundles.searchIndex.OpenSearchBulkSink; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; @@ -125,10 +127,10 @@ public static boolean isCoordinatingJob(UUID jobId) { private IndexingFailureRecorder failureRecorder; private BulkSink searchIndexSink; - // Per-entity staged index promotion + // Per-entity index promotion private EntityCompletionTracker entityTracker; - private RecreateIndexHandler indexPromotionHandler; - private ReindexContext stagedIndexContext; + private RecreateIndexHandler recreateIndexHandler; + private ReindexContext recreateContext; // Reader stats tracking (accumulated across all worker threads) private final AtomicLong coordinatorReaderSuccess = new AtomicLong(0); @@ -193,8 +195,8 @@ public void setAppContext(UUID appId, Long startTime) { } /** - * Set the job notifier for alerting other servers when a job starts. Servers discover the job - * through database polling. + * Set the job notifier for alerting other servers when a job starts. When set, other servers in + * the cluster will be notified via Redis Pub/Sub (if available) or discovered via polling. * * @param notifier The job notifier */ @@ -309,19 +311,19 @@ public Optional joinJob(UUID jobId) { * none remain 3. Coordinates with other servers for load balancing * * @param bulkSink The sink for writing to search index - * @param stagedIndexContext Context for staged index writes and promotion + * @param recreateContext Context for index recreation, if applicable + * @param recreateIndex Whether indices should be recreated * @return Execution result with statistics */ public ExecutionResult execute( - BulkSink bulkSink, ReindexContext stagedIndexContext, ReindexingConfiguration reindexConfig) { + BulkSink bulkSink, + ReindexContext recreateContext, + boolean recreateIndex, + ReindexingConfiguration reindexConfig) { if (currentJob == null) { throw new IllegalStateException("No job to execute - call createJob() or joinJob() first"); } - if (stagedIndexContext == null || stagedIndexContext.isEmpty()) { - throw new IllegalArgumentException( - "Staged index context is required for distributed reindexing"); - } UUID jobId = currentJob.getId(); LOG.info("Server {} starting execution of job {}", serverId, jobId); @@ -404,12 +406,12 @@ public ExecutionResult execute( // Stats are tracked per-entityType by StageStatsTracker in PartitionWorker // No need for redundant server-level stats persistence - // Store staged index context for per-entity promotion - this.stagedIndexContext = stagedIndexContext; + // Store recreate context for per-entity promotion + this.recreateContext = recreateContext; // Initialize entity completion tracker for per-entity index promotion this.entityTracker = new EntityCompletionTracker(jobId); - initializeEntityTracker(jobId); + initializeEntityTracker(jobId, recreateIndex); coordinator.setEntityCompletionTracker(entityTracker); // Start lock refresh thread to prevent lock expiration during long-running jobs @@ -460,7 +462,8 @@ public ExecutionResult execute( workerId, bulkSink, batchSize, - stagedIndexContext, + recreateContext, + recreateIndex, totalSuccess, totalFailed, reindexConfig); @@ -488,7 +491,7 @@ public ExecutionResult execute( // Final reconciliation pass: catch ALL participant-server completions before // the stale-reclaimer is killed. Participant workers may have finished partitions // that were never reconciled by the stale-reclaimer's periodic loop. - if (entityTracker != null && stagedIndexContext != null) { + if (entityTracker != null && recreateContext != null) { LOG.info("Running final DB reconciliation for job {}", jobId); List allPartitions = coordinator.getPartitions(jobId, null); entityTracker.reconcileFromDatabase(allPartitions); @@ -653,7 +656,8 @@ private void runWorkerLoop( int workerId, BulkSink bulkSink, int batchSize, - ReindexContext stagedIndexContext, + ReindexContext recreateContext, + boolean recreateIndex, AtomicLong totalSuccess, AtomicLong totalFailed, ReindexingConfiguration reindexConfig) { @@ -662,7 +666,13 @@ private void runWorkerLoop( PartitionWorker worker = new PartitionWorker( - coordinator, bulkSink, batchSize, stagedIndexContext, failureRecorder, reindexConfig); + coordinator, + bulkSink, + batchSize, + recreateContext, + recreateIndex, + failureRecorder, + reindexConfig); synchronized (activeWorkers) { activeWorkers.add(worker); @@ -1070,7 +1080,7 @@ public void updateStagedIndexMapping(Map stagedIndexMapping) { /** * Initialize the entity completion tracker with partition counts and promotion callback. */ - private void initializeEntityTracker(UUID jobId) { + private void initializeEntityTracker(UUID jobId, boolean recreateIndex) { // Count partitions per entity Map partitionCountByEntity = new HashMap<>(); List allPartitions = coordinator.getPartitions(jobId, null); @@ -1089,60 +1099,79 @@ private void initializeEntityTracker(UUID jobId) { partitionCountByEntity.size(), partitionCountByEntity); - if (partitionCountByEntity.isEmpty()) { - LOG.info("No partitions found for job {}; finalizer will promote staged indexes", jobId); - return; - } - - if (stagedIndexContext == null || stagedIndexContext.isEmpty()) { - throw new IllegalStateException("Staged index context is required for entity promotion"); - } - indexPromotionHandler = Entity.getSearchRepository().createReindexHandler(); - // Wire job configuration so applyLiveServingSettings can revert bulk-build overrides - // (refresh=-1, replicas=0, async translog) before the per-entity alias swap. - if (indexPromotionHandler instanceof DefaultRecreateHandler defaultHandler - && currentJob != null - && currentJob.getJobConfiguration() != null) { - defaultHandler.withJobData(currentJob.getJobConfiguration()); + // Set up per-entity promotion callback if recreating indices + if (recreateIndex && recreateContext != null) { + this.recreateIndexHandler = Entity.getSearchRepository().createReindexHandler(); + // Wire jobData into the handler so applyLiveServingSettings can revert bulk-build + // overrides (refresh_interval=-1, replicas=0, async translog) before the per-entity + // alias swap. Without this, buildRevertJson returns null and the bulk overrides + // silently become the live settings. + if (recreateIndexHandler instanceof DefaultRecreateHandler defaultHandler + && currentJob != null + && currentJob.getJobConfiguration() != null) { + defaultHandler.withJobData(currentJob.getJobConfiguration()); + } + entityTracker.setOnEntityComplete(this::promoteEntityIndex); + LOG.info( + "Per-entity promotion callback SET for job {} (recreateIndex={}, recreateContext entities={})", + jobId, + recreateIndex, + recreateContext.getEntities()); + } else { + LOG.info( + "Per-entity promotion callback NOT set for job {} (recreateIndex={}, recreateContext={})", + jobId, + recreateIndex, + recreateContext != null ? "present" : "null"); } - entityTracker.setOnEntityComplete(this::promoteEntityIndex); - LOG.info( - "Per-entity promotion callback set for job {} (staged index entities={})", - jobId, - stagedIndexContext.getEntities()); } /** * Promote a single entity's index when all its partitions complete. */ private void promoteEntityIndex(String entityType, boolean success) { - if (indexPromotionHandler == null || stagedIndexContext == null) { + if (recreateIndexHandler == null || recreateContext == null) { LOG.warn( - "Cannot promote index for entity '{}' - no index promotion handler or staged context", + "Cannot promote index for entity '{}' - no recreateIndexHandler or recreateContext", entityType); return; } - EntityReindexContext entityContext = - EntityReindexContextMapper.fromStagedContext(stagedIndexContext, entityType); - if (entityContext.getStagedIndex() == null) { + Optional stagedIndexOpt = recreateContext.getStagedIndex(entityType); + if (stagedIndexOpt.isEmpty()) { LOG.debug("No staged index for entity '{}', skipping promotion", entityType); return; } try { + String canonicalIndex = recreateContext.getCanonicalIndex(entityType).orElse(null); + String originalIndex = recreateContext.getOriginalIndex(entityType).orElse(null); + LOG.debug( "Promoting entity '{}': success={}, canonicalIndex={}, stagedIndex={}", entityType, success, - entityContext.getCanonicalIndex(), - entityContext.getStagedIndex()); - - if (indexPromotionHandler instanceof DefaultRecreateHandler defaultHandler) { + canonicalIndex, + stagedIndexOpt.get()); + + EntityReindexContext entityContext = + EntityReindexContext.builder() + .entityType(entityType) + .originalIndex(originalIndex) + .canonicalIndex(canonicalIndex) + .activeIndex(originalIndex) + .stagedIndex(stagedIndexOpt.get()) + .canonicalAliases(recreateContext.getCanonicalAlias(entityType).orElse(null)) + .existingAliases(recreateContext.getExistingAliases(entityType)) + .parentAliases( + new HashSet<>(listOrEmpty(recreateContext.getParentAliases(entityType)))) + .build(); + + if (recreateIndexHandler instanceof DefaultRecreateHandler defaultHandler) { LOG.info("Promoting index for entity '{}' (success={})", entityType, success); defaultHandler.promoteEntityIndex(entityContext, success); } else { - indexPromotionHandler.finalizeReindex(entityContext, success); + recreateIndexHandler.finalizeReindex(entityContext, success); } } catch (Exception e) { LOG.error("Failed to promote index for entity '{}'", entityType, e); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java index 5a50b9c86422..3079b00aa801 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionCalculator.java @@ -24,7 +24,6 @@ import org.openmetadata.service.Entity; import org.openmetadata.service.apps.bundles.searchIndex.EntityPriority; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; -import org.openmetadata.service.apps.bundles.searchIndex.SearchIndexEntityTypes; import org.openmetadata.service.jdbi3.EntityRepository; import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; import org.openmetadata.service.jdbi3.ListFilter; @@ -82,6 +81,18 @@ public class PartitionCalculator { Map.entry("queryCostRecord", 0.3) // Time series, simple structure ); + /** Time series entity types */ + private static final Set TIME_SERIES_ENTITIES = + Set.of( + "testCaseResolutionStatus", + "testCaseResult", + "queryCostRecord", + "webAnalyticEntityViewReportData", + "webAnalyticUserActivityReportData", + "entityReportData", + "rawCostAnalysisReportData", + "aggregatedCostAnalysisReportData"); + private final int partitionSize; private final int minPartitionsPerEntity; @@ -245,7 +256,7 @@ public long getEntityCount(String entityType) { public long getEntityCount(String entityType, ReindexingConfiguration reindexConfig) { try { long count; - if (SearchIndexEntityTypes.isTimeSeriesEntity(entityType)) { + if (TIME_SERIES_ENTITIES.contains(entityType)) { count = getTimeSeriesEntityCount(entityType, reindexConfig); } else { count = getRegularEntityCount(entityType); @@ -267,7 +278,7 @@ private long getTimeSeriesEntityCount(String entityType, ReindexingConfiguration ListFilter listFilter = new ListFilter(Include.ALL); EntityTimeSeriesRepository repository; - if (SearchIndexEntityTypes.isDataInsightEntity(entityType)) { + if (isDataInsightIndex(entityType)) { listFilter.addQueryParam("entityFQNHash", FullyQualifiedName.buildHash(entityType)); repository = Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA); } else { @@ -292,6 +303,10 @@ private long getTimeSeriesEntityCount(String entityType, ReindexingConfiguration return repository.getTimeSeriesDao().listCount(listFilter); } + private boolean isDataInsightIndex(String entityType) { + return entityType.endsWith("ReportData"); + } + /** * Get entity counts for all requested entity types. * diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java index c108a4ba0094..1db12e6c577d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorker.java @@ -14,9 +14,13 @@ package org.openmetadata.service.apps.bundles.searchIndex.distributed; import static org.openmetadata.common.utils.CommonUtil.listOrEmpty; +import static org.openmetadata.service.Entity.QUERY_COST_RECORD; +import static org.openmetadata.service.Entity.TEST_CASE_RESOLUTION_STATUS; +import static org.openmetadata.service.Entity.TEST_CASE_RESULT; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @@ -24,6 +28,7 @@ import org.apache.commons.lang3.exception.ExceptionUtils; import org.openmetadata.schema.EntityInterface; import org.openmetadata.schema.EntityTimeSeriesInterface; +import org.openmetadata.schema.analytics.ReportData; import org.openmetadata.schema.system.EntityError; import org.openmetadata.schema.type.Include; import org.openmetadata.schema.utils.ResultList; @@ -31,7 +36,6 @@ import org.openmetadata.service.apps.bundles.searchIndex.BulkSink; import org.openmetadata.service.apps.bundles.searchIndex.IndexingFailureRecorder; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; -import org.openmetadata.service.apps.bundles.searchIndex.SearchIndexEntityTypes; import org.openmetadata.service.apps.bundles.searchIndex.stats.StageStatsTracker; import org.openmetadata.service.cache.EntityCacheBypass; import org.openmetadata.service.exception.SearchIndexException; @@ -52,14 +56,26 @@ public class PartitionWorker { private static final long MAX_CURSOR_INITIALIZATION_OFFSET = (long) Integer.MAX_VALUE + 1L; + /** Time series entity types that need special handling */ + static final Set TIME_SERIES_ENTITIES = + Set.of( + ReportData.ReportDataType.ENTITY_REPORT_DATA.value(), + ReportData.ReportDataType.RAW_COST_ANALYSIS_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(), + ReportData.ReportDataType.WEB_ANALYTIC_ENTITY_VIEW_REPORT_DATA.value(), + ReportData.ReportDataType.AGGREGATED_COST_ANALYSIS_REPORT_DATA.value(), + TEST_CASE_RESOLUTION_STATUS, + TEST_CASE_RESULT, + QUERY_COST_RECORD); + /** Context key for entity type */ private static final String ENTITY_TYPE_KEY = "entityType"; - /** Context key used by search sinks to write into staged indexes. */ - private static final String STAGED_WRITE_KEY = "recreateIndex"; + /** Context key for recreate index flag */ + private static final String RECREATE_INDEX = "recreateIndex"; - /** Context key for staged index context. */ - private static final String STAGED_CONTEXT_KEY = "recreateContext"; + /** Context key for recreate context */ + private static final String RECREATE_CONTEXT = "recreateContext"; /** Context key for target index */ private static final String TARGET_INDEX_KEY = "targetIndex"; @@ -76,7 +92,8 @@ public class PartitionWorker { private final DistributedSearchIndexCoordinator coordinator; private final BulkSink searchIndexSink; private final int batchSize; - private final ReindexContext stagedIndexContext; + private final ReindexContext recreateContext; + private final boolean recreateIndex; private final AtomicBoolean stopped = new AtomicBoolean(false); private final IndexingFailureRecorder failureRecorder; private final ReindexingConfiguration reindexConfig; @@ -85,30 +102,41 @@ public PartitionWorker( DistributedSearchIndexCoordinator coordinator, BulkSink searchIndexSink, int batchSize, - ReindexContext stagedIndexContext) { - this(coordinator, searchIndexSink, batchSize, stagedIndexContext, null, null); + ReindexContext recreateContext, + boolean recreateIndex) { + this(coordinator, searchIndexSink, batchSize, recreateContext, recreateIndex, null, null); } public PartitionWorker( DistributedSearchIndexCoordinator coordinator, BulkSink searchIndexSink, int batchSize, - ReindexContext stagedIndexContext, + ReindexContext recreateContext, + boolean recreateIndex, IndexingFailureRecorder failureRecorder) { - this(coordinator, searchIndexSink, batchSize, stagedIndexContext, failureRecorder, null); + this( + coordinator, + searchIndexSink, + batchSize, + recreateContext, + recreateIndex, + failureRecorder, + null); } public PartitionWorker( DistributedSearchIndexCoordinator coordinator, BulkSink searchIndexSink, int batchSize, - ReindexContext stagedIndexContext, + ReindexContext recreateContext, + boolean recreateIndex, IndexingFailureRecorder failureRecorder, ReindexingConfiguration reindexConfig) { this.coordinator = coordinator; this.searchIndexSink = searchIndexSink; this.batchSize = batchSize; - this.stagedIndexContext = stagedIndexContext; + this.recreateContext = recreateContext; + this.recreateIndex = recreateIndex; this.failureRecorder = failureRecorder; this.reindexConfig = reindexConfig; } @@ -133,7 +161,7 @@ public PartitionResult processPartition(SearchIndexPartition partition) { } private PartitionResult processPartitionInternal(SearchIndexPartition partition) { - String entityType = SearchIndexEntityTypes.normalizeEntityType(partition.getEntityType()); + String entityType = partition.getEntityType(); long rangeStart = partition.getRangeStart(); long rangeEnd = partition.getRangeEnd(); @@ -581,21 +609,22 @@ private void recordReaderFailures( */ private ResultList readEntitiesKeyset(String entityType, String keysetCursor, int limit) throws SearchIndexException { - String normalizedEntityType = SearchIndexEntityTypes.normalizeEntityType(entityType); - // Selective fields avoid running expensive field fetchers that are stripped out before - // indexing. - List fields = ReindexingUtil.getSearchIndexFields(normalizedEntityType); + // Selective fields, not "*". Asking for "*" runs every registered fieldFetcher in + // setFieldsInBulk — including expensive ones like fetchAndSetOwns on Team/User where every + // owned entity becomes an Entity.getEntityReferenceById round-trip — and the index class then + // strips most of those out via getExcludedFields anyway. Mirror what EntityReader does on the + // single-server pipeline (PR #27723) so both paths request the same minimal set. + List fields = ReindexingUtil.getSearchIndexFields(entityType); - if (!SearchIndexEntityTypes.isTimeSeriesEntity(normalizedEntityType)) { - PaginatedEntitiesSource source = - new PaginatedEntitiesSource(normalizedEntityType, limit, fields, 0); + if (!TIME_SERIES_ENTITIES.contains(entityType)) { + PaginatedEntitiesSource source = new PaginatedEntitiesSource(entityType, limit, fields, 0); return source.readNextKeyset(keysetCursor); } else { Long filterStartTs = null; Long filterEndTs = null; if (reindexConfig != null) { - long startTs = reindexConfig.getTimeSeriesStartTs(normalizedEntityType); + long startTs = reindexConfig.getTimeSeriesStartTs(entityType); if (startTs > 0) { filterStartTs = startTs; filterEndTs = System.currentTimeMillis(); @@ -604,8 +633,8 @@ private ResultList readEntitiesKeyset(String entityType, String keysetCursor, PaginatedEntityTimeSeriesSource source = (filterStartTs != null) ? new PaginatedEntityTimeSeriesSource( - normalizedEntityType, limit, fields, filterStartTs, filterEndTs) - : new PaginatedEntityTimeSeriesSource(normalizedEntityType, limit, fields, 0); + entityType, limit, fields, filterStartTs, filterEndTs) + : new PaginatedEntityTimeSeriesSource(entityType, limit, fields, 0); return source.readWithCursor(keysetCursor); } } @@ -614,8 +643,8 @@ private String initializeKeysetCursor(SearchIndexPartition partition, long offse if (offset <= 0) { return null; } - String entityType = SearchIndexEntityTypes.normalizeEntityType(partition.getEntityType()); - if (SearchIndexEntityTypes.isTimeSeriesEntity(entityType)) { + String entityType = partition.getEntityType(); + if (TIME_SERIES_ENTITIES.contains(entityType)) { return RestUtil.encodeCursor(String.valueOf(offset)); } // Fast path: coordinator precomputed boundary cursors for every partition's @@ -664,9 +693,8 @@ private int toCursorOffset(String entityType, long offset) { private void writeToSink( String entityType, ResultList resultList, Map contextData) throws Exception { - String normalizedEntityType = SearchIndexEntityTypes.normalizeEntityType(entityType); - if (!SearchIndexEntityTypes.isTimeSeriesEntity(normalizedEntityType)) { + if (!TIME_SERIES_ENTITIES.contains(entityType)) { List entities = (List) resultList.getData(); searchIndexSink.write(entities, contextData); } else { @@ -684,30 +712,21 @@ private void writeToSink( * @return Context data map */ private Map createContextData(String entityType, StageStatsTracker statsTracker) { - String normalizedEntityType = SearchIndexEntityTypes.normalizeEntityType(entityType); Map contextData = new java.util.HashMap<>(); - contextData.put(ENTITY_TYPE_KEY, normalizedEntityType); - contextData.put(STAGED_WRITE_KEY, true); + contextData.put(ENTITY_TYPE_KEY, entityType); + contextData.put(RECREATE_INDEX, recreateIndex); if (statsTracker != null) { contextData.put(BulkSink.STATS_TRACKER_CONTEXT_KEY, statsTracker); } - if (stagedIndexContext == null) { - throw new IllegalStateException( - "Staged index context is required for distributed reindexing"); + if (recreateContext != null) { + contextData.put(RECREATE_CONTEXT, recreateContext); + recreateContext + .getStagedIndex(entityType) + .ifPresent(index -> contextData.put(TARGET_INDEX_KEY, index)); } - String targetIndex = - stagedIndexContext - .getStagedIndex(normalizedEntityType) - .orElseThrow( - () -> - new IllegalStateException( - "No staged index configured for entity type: " + normalizedEntityType)); - contextData.put(STAGED_CONTEXT_KEY, stagedIndexContext); - contextData.put(TARGET_INDEX_KEY, targetIndex); - return contextData; } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifier.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifier.java index 31f643cf71ec..0d41ecf0f47d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifier.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifier.java @@ -13,7 +13,6 @@ package org.openmetadata.service.apps.bundles.searchIndex.distributed; -import java.util.HashSet; import java.util.List; import java.util.Set; import java.util.UUID; @@ -27,31 +26,28 @@ import org.openmetadata.service.jdbi3.CollectionDAO; /** - * Database polling based job notifier for distributed job discovery. + * Database polling based job notifier as fallback when Redis is not available. * *

Uses adaptive polling intervals: * *

    - *
  • 1 second while actively participating in a job - *
  • 2 seconds plus jitter while recently started or after job activity - *
  • 30 seconds plus jitter after an extended idle period + *
  • 30 seconds when idle (no active jobs) + *
  • 1 second when actively participating in a job *
+ * + *

This minimizes database overhead while still providing reasonable job discovery latency. */ @Slf4j public class PollingJobNotifier implements DistributedJobNotifier { - private static final long FAST_IDLE_POLL_INTERVAL_MS = 2_000; - private static final long BACKOFF_IDLE_POLL_INTERVAL_MS = 30_000; + /** Poll interval when no job is running (30 seconds) */ + private static final long IDLE_POLL_INTERVAL_MS = 30_000; + /** Poll interval when actively participating (1 second) */ private static final long ACTIVE_POLL_INTERVAL_MS = 1_000; - private static final long FAST_IDLE_WINDOW_MS = 60_000; - private static final long FAST_IDLE_JITTER_MS = 1_000; - private static final long BACKOFF_IDLE_JITTER_MS = 5_000; private final CollectionDAO collectionDAO; private final String serverId; - private final long fastIdleJitterMs; - private final long backoffIdleJitterMs; private final AtomicBoolean running = new AtomicBoolean(false); private final AtomicBoolean participating = new AtomicBoolean(false); private final Set knownJobs = ConcurrentHashMap.newKeySet(); @@ -59,13 +55,10 @@ public class PollingJobNotifier implements DistributedJobNotifier { private ScheduledExecutorService scheduler; private Consumer jobStartedCallback; private volatile long lastPollTime = 0; - private volatile long fastIdleUntil = 0; public PollingJobNotifier(CollectionDAO collectionDAO, String serverId) { this.collectionDAO = collectionDAO; this.serverId = serverId; - this.fastIdleJitterMs = computeJitter(FAST_IDLE_JITTER_MS, 17); - this.backoffIdleJitterMs = computeJitter(BACKOFF_IDLE_JITTER_MS, 31); } @Override @@ -75,10 +68,6 @@ public void start() { return; } - long now = System.currentTimeMillis(); - lastPollTime = 0; - extendFastIdleWindow(now); - scheduler = Executors.newSingleThreadScheduledExecutor( Thread.ofPlatform() @@ -86,14 +75,14 @@ public void start() { "reindex-job-notifier-" + serverId.substring(0, Math.min(8, serverId.length()))) .factory()); + // Schedule with fixed delay of 1 second, but actual polling is controlled by interval logic scheduler.scheduleWithFixedDelay( this::pollForJobs, 0, ACTIVE_POLL_INTERVAL_MS, TimeUnit.MILLISECONDS); LOG.info( - "PollingJobNotifier started on server {} (fast idle: {}s, backoff idle: {}s, active: {}s)", + "PollingJobNotifier started on server {} (idle: {}s, active: {}s)", serverId, - FAST_IDLE_POLL_INTERVAL_MS / 1000, - BACKOFF_IDLE_POLL_INTERVAL_MS / 1000, + IDLE_POLL_INTERVAL_MS / 1000, ACTIVE_POLL_INTERVAL_MS / 1000); } @@ -121,8 +110,9 @@ public void stop() { @Override public void notifyJobStarted(UUID jobId, String jobType) { + // In polling mode, we don't actively notify - other servers will discover via polling + // But we track it locally to avoid re-notifying ourselves knownJobs.add(jobId); - extendFastIdleWindow(System.currentTimeMillis()); LOG.debug( "Job {} (type: {}) started - other servers will discover via polling", jobId, jobType); } @@ -130,7 +120,6 @@ public void notifyJobStarted(UUID jobId, String jobType) { @Override public void notifyJobCompleted(UUID jobId) { knownJobs.remove(jobId); - extendFastIdleWindow(System.currentTimeMillis()); LOG.debug("Job {} completed - removed from known jobs", jobId); } @@ -155,9 +144,6 @@ public String getType() { */ public void setParticipating(boolean isParticipating) { this.participating.set(isParticipating); - if (!isParticipating) { - extendFastIdleWindow(System.currentTimeMillis()); - } } private void pollForJobs() { @@ -166,23 +152,32 @@ private void pollForJobs() { } long now = System.currentTimeMillis(); - if (now - lastPollTime < currentPollIntervalMs(now)) { + long interval = participating.get() ? ACTIVE_POLL_INTERVAL_MS : IDLE_POLL_INTERVAL_MS; + + // Skip poll if not enough time has elapsed + if (now - lastPollTime < interval) { return; } lastPollTime = now; try { + // Fast, lightweight query for running jobs List runningJobIds = collectionDAO.searchIndexJobDAO().getRunningJobIds(); if (runningJobIds.isEmpty()) { - handleNoRunningJobs(now); + // No jobs running - clear known jobs and stay in idle mode + if (!knownJobs.isEmpty()) { + LOG.debug("No running jobs found, clearing {} known jobs", knownJobs.size()); + knownJobs.clear(); + } return; } - extendFastIdleWindow(now); + // Check for new jobs we haven't seen for (String jobIdStr : runningJobIds) { UUID jobId = UUID.fromString(jobIdStr); if (!knownJobs.contains(jobId)) { + // New job discovered! LOG.info("Discovered new running job via polling: {}", jobId); knownJobs.add(jobId); @@ -192,38 +187,12 @@ private void pollForJobs() { } } - Set runningJobIdSet = new HashSet<>(runningJobIds); - knownJobs.removeIf(jobId -> !runningJobIdSet.contains(jobId.toString())); + // Clean up jobs that are no longer running + knownJobs.removeIf( + jobId -> runningJobIds.stream().noneMatch(id -> id.equals(jobId.toString()))); } catch (Exception e) { LOG.error("Error polling for jobs", e); } } - - private void handleNoRunningJobs(long now) { - if (knownJobs.isEmpty()) { - return; - } - LOG.debug("No running jobs found, clearing {} known jobs", knownJobs.size()); - knownJobs.clear(); - extendFastIdleWindow(now); - } - - private long currentPollIntervalMs(long now) { - if (participating.get()) { - return ACTIVE_POLL_INTERVAL_MS; - } - if (now <= fastIdleUntil) { - return FAST_IDLE_POLL_INTERVAL_MS + fastIdleJitterMs; - } - return BACKOFF_IDLE_POLL_INTERVAL_MS + backoffIdleJitterMs; - } - - private void extendFastIdleWindow(long now) { - fastIdleUntil = now + FAST_IDLE_WINDOW_MS; - } - - private long computeJitter(long maxJitterMs, int salt) { - return Math.floorMod((serverId.hashCode() * 31) + salt, (int) maxJitterMs + 1); - } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/RedisJobNotifier.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/RedisJobNotifier.java new file mode 100644 index 000000000000..0d163ef79283 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/RedisJobNotifier.java @@ -0,0 +1,245 @@ +/* + * Copyright 2024 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.openmetadata.service.apps.bundles.searchIndex.distributed; + +import io.lettuce.core.RedisClient; +import io.lettuce.core.RedisURI; +import io.lettuce.core.api.StatefulRedisConnection; +import io.lettuce.core.pubsub.RedisPubSubAdapter; +import io.lettuce.core.pubsub.StatefulRedisPubSubConnection; +import java.time.Duration; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.service.cache.CacheConfig; + +/** + * Redis Pub/Sub based job notifier for instant push notifications. + * + *

When Redis is available, this provides zero-latency job discovery across all servers in the + * cluster. Messages are delivered instantly via Redis Pub/Sub. + */ +@Slf4j +public class RedisJobNotifier implements DistributedJobNotifier { + + private static final String CHANNEL_PREFIX = "om:distributed-jobs:"; + private static final String START_CHANNEL = CHANNEL_PREFIX + "start"; + private static final String COMPLETE_CHANNEL = CHANNEL_PREFIX + "complete"; + + private final CacheConfig.Redis redisConfig; + private final String serverId; + private final AtomicBoolean running = new AtomicBoolean(false); + + private RedisClient redisClient; + private StatefulRedisPubSubConnection subConnection; + private StatefulRedisConnection pubConnection; + private Consumer jobStartedCallback; + + public RedisJobNotifier(CacheConfig cacheConfig, String serverId) { + this.redisConfig = cacheConfig.redis; + this.serverId = serverId; + } + + @Override + public void start() { + if (!running.compareAndSet(false, true)) { + LOG.warn("RedisJobNotifier already running"); + return; + } + + try { + RedisURI uri = buildRedisURI(); + redisClient = RedisClient.create(uri); + + // Create subscription connection + subConnection = redisClient.connectPubSub(); + subConnection.addListener( + new RedisPubSubAdapter<>() { + @Override + public void message(String channel, String message) { + handleMessage(channel, message); + } + }); + + // Subscribe to job channels + subConnection.sync().subscribe(START_CHANNEL, COMPLETE_CHANNEL); + + // Create publish connection (separate from subscription) + pubConnection = redisClient.connect(); + + LOG.info( + "RedisJobNotifier started on server {} - subscribed to channels: {}, {}", + serverId, + START_CHANNEL, + COMPLETE_CHANNEL); + + } catch (Exception e) { + running.set(false); + LOG.error("Failed to start RedisJobNotifier", e); + throw new RuntimeException("Failed to initialize Redis Pub/Sub", e); + } + } + + @Override + public void stop() { + if (!running.compareAndSet(true, false)) { + return; + } + + try { + if (subConnection != null) { + subConnection.sync().unsubscribe(START_CHANNEL, COMPLETE_CHANNEL); + subConnection.close(); + } + if (pubConnection != null) { + pubConnection.close(); + } + if (redisClient != null) { + redisClient.shutdown(); + } + LOG.info("RedisJobNotifier stopped on server {}", serverId); + } catch (Exception e) { + LOG.error("Error stopping RedisJobNotifier", e); + } + } + + @Override + public void notifyJobStarted(UUID jobId, String jobType) { + if (!running.get() || pubConnection == null) { + LOG.warn("Cannot notify job started - RedisJobNotifier not running"); + return; + } + + try { + String message = formatMessage(jobId, jobType, serverId); + long receivers = pubConnection.sync().publish(START_CHANNEL, message); + LOG.info( + "Published job start notification for {} (type: {}) to {} subscribers", + jobId, + jobType, + receivers); + } catch (Exception e) { + LOG.error("Failed to publish job start notification for {}", jobId, e); + } + } + + @Override + public void notifyJobCompleted(UUID jobId) { + if (!running.get() || pubConnection == null) { + LOG.warn("Cannot notify job completed - RedisJobNotifier not running"); + return; + } + + try { + String message = formatMessage(jobId, "COMPLETED", serverId); + pubConnection.sync().publish(COMPLETE_CHANNEL, message); + LOG.debug("Published job completion notification for {}", jobId); + } catch (Exception e) { + LOG.error("Failed to publish job completion notification for {}", jobId, e); + } + } + + @Override + public void onJobStarted(Consumer callback) { + this.jobStartedCallback = callback; + } + + @Override + public boolean isRunning() { + return running.get(); + } + + @Override + public String getType() { + return "redis-pubsub"; + } + + private void handleMessage(String channel, String message) { + try { + String[] parts = message.split("\\|"); + if (parts.length < 3) { + LOG.warn("Invalid message format: {}", message); + return; + } + + UUID jobId = UUID.fromString(parts[0]); + String jobType = parts[1]; + String sourceServer = parts[2]; + + // Don't process our own messages + if (serverId.equals(sourceServer)) { + LOG.debug("Ignoring own message for job {}", jobId); + return; + } + + if (START_CHANNEL.equals(channel)) { + LOG.info( + "Received job start notification from server {}: job={}, type={}", + sourceServer, + jobId, + jobType); + if (jobStartedCallback != null) { + jobStartedCallback.accept(jobId); + } + } else if (COMPLETE_CHANNEL.equals(channel)) { + LOG.debug("Received job completion notification: job={}", jobId); + } + + } catch (Exception e) { + LOG.error("Error handling message on channel {}: {}", channel, message, e); + } + } + + private String formatMessage(UUID jobId, String jobType, String sourceServer) { + return jobId.toString() + "|" + jobType + "|" + sourceServer; + } + + private RedisURI buildRedisURI() { + String url = redisConfig.url; + RedisURI.Builder builder; + + if (url.startsWith("redis://") || url.startsWith("rediss://")) { + RedisURI uri = RedisURI.create(url); + builder = + RedisURI.Builder.redis(uri.getHost(), uri.getPort()) + .withTimeout(Duration.ofMillis(redisConfig.connectTimeoutMs)); + } else if (url.contains(":")) { + String[] parts = url.split(":"); + String host = parts[0]; + int port = Integer.parseInt(parts[1]); + builder = + RedisURI.Builder.redis(host, port) + .withTimeout(Duration.ofMillis(redisConfig.connectTimeoutMs)); + } else { + builder = + RedisURI.Builder.redis(url).withTimeout(Duration.ofMillis(redisConfig.connectTimeoutMs)); + } + + if (redisConfig.authType == CacheConfig.AuthType.PASSWORD) { + if (redisConfig.username != null && redisConfig.passwordRef != null) { + builder.withAuthentication(redisConfig.username, redisConfig.passwordRef); + } else if (redisConfig.passwordRef != null) { + builder.withPassword(redisConfig.passwordRef.toCharArray()); + } + } + + if (redisConfig.useSSL) { + builder.withSsl(true); + } + + builder.withDatabase(redisConfig.database); + return builder.build(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java index 0de30434db7b..cc9a3f8384d9 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListener.java @@ -27,7 +27,10 @@ public LoggingProgressListener() { @Override public void onJobStarted(ReindexingJobContext context) { LOG.info( - "Reindexing job started - Job ID: {}, Source: {}", context.getJobId(), context.getSource()); + "Reindexing job started - Job ID: {}, Source: {}, Distributed: {}", + context.getJobId(), + context.getSource(), + context.isDistributed()); } @Override @@ -42,16 +45,17 @@ public void onJobConfigured(ReindexingJobContext context, ReindexingConfiguratio logger.addInitDetail("Max Concurrent Requests", config.maxConcurrentRequests()); logger.addInitDetail("Payload Size", formatBytes(config.payloadSize())); logger.addInitDetail("Auto-tune", config.autoTune() ? "Enabled" : "Disabled"); - logger.addInitDetail("Indexing Mode", "Staged indexes with alias promotion"); + logger.addInitDetail("Recreate Index", config.recreateIndex() ? "Yes" : "No"); + logger.addInitDetail("Distributed Mode", config.useDistributedIndexing() ? "Yes" : "No"); logger.logInitialization(); } @Override public void onIndexRecreationStarted(Set entities) { - LOG.info("Preparing staged indexes for {} entity types", entities.size()); + LOG.info("Starting index recreation for {} entity types", entities.size()); if (LOG.isDebugEnabled()) { - LOG.debug("Entities to stage: {}", String.join(", ", entities)); + LOG.debug("Entities to recreate: {}", String.join(", ", entities)); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListener.java b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListener.java index 0a1a2a9131ba..c3f70b5defb9 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListener.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListener.java @@ -27,8 +27,7 @@ public class SlackProgressListener implements ReindexingProgressListener { private static final String PRODUCER_THREADS = "Producer threads"; private static final String TOTAL_ENTITIES = "Total entities"; private static final String QUEUE_SIZE = "Queue size"; - private static final String INDEXING_MODE = "Indexing mode"; - private static final String STAGED_PROMOTION = "Staged indexes with alias promotion"; + private static final String RECREATING_INDICES = "Recreating indices"; private static final String PAYLOAD_SIZE = "Payload size"; private static final String CONCURRENT_REQUESTS = "Concurrent requests"; @@ -59,8 +58,7 @@ public void onJobConfigured(ReindexingJobContext context, ReindexingConfiguratio @Override public void onIndexRecreationStarted(Set entities) { - LOG.debug( - "Slack notification: Staged index preparation started for {} entities", entities.size()); + LOG.debug("Slack notification: Index recreation started for {} entities", entities.size()); } @Override @@ -127,7 +125,7 @@ private Map buildConfigDetails(ReindexingConfiguration config) { details.put(PRODUCER_THREADS, String.valueOf(config.producerThreads())); details.put(QUEUE_SIZE, String.valueOf(config.queueSize())); details.put(TOTAL_ENTITIES, String.valueOf(totalEntities)); - details.put(INDEXING_MODE, STAGED_PROMOTION); + details.put(RECREATING_INDICES, config.recreateIndex() ? "Yes" : "No"); details.put(PAYLOAD_SIZE, (config.payloadSize() / (1024 * 1024)) + " MB"); details.put(CONCURRENT_REQUESTS, String.valueOf(config.maxConcurrentRequests())); return details; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/AppRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/AppRepository.java index e730e8d89c9d..3c965f107494 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/AppRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/AppRepository.java @@ -397,43 +397,6 @@ public ResultList listAppExtensionAfterTimeByName( } } - /** - * Page through extensions inside a half-open {@code [startTime, endTime)} window. Unlike - * {@link #listAppExtensionAfterTimeByName}, the SQL filter excludes rows at or after - * {@code endTime} so OFFSET pagination stays correct even when new rows are inserted - * concurrently. Useful for any counter that aggregates across multiple pages. - * - *

Known limitation: the {@code apps_extension_time_series} table has no surrogate - * primary key, so the ORDER BY tie-breaker is limited to {@code timestamp}. Two writes that - * land in the same millisecond can be ordered non-deterministically across separate page - * queries, causing one row near a page boundary to be skipped or counted twice. At - * {@link org.openmetadata.service.resources.mcp.McpUsageResource}'s page size (1000) the risk - * is bounded and acceptable for a growth-metric dashboard. Adding a deterministic - * tie-breaker would require a schema migration to introduce a surrogate id column. - */ - public List listAppExtensionInWindowByName( - App app, - long startTime, - long endTime, - int limitParam, - int offset, - Class clazz, - AppExtension.ExtensionType extensionType) { - if (limitParam <= 0) { - return new ArrayList<>(); - } - List jsons = - daoCollection - .appExtensionTimeSeriesDao() - .listAppExtensionInWindowByName( - app.getName(), limitParam, offset, startTime, endTime, extensionType.toString()); - List entities = new ArrayList<>(jsons.size()); - for (String json : jsons) { - entities.add(JsonUtils.readValue(json, clazz)); - } - return entities; - } - public ResultList listAppExtensionAfterTimeById( App app, long startTime, diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java index 56f1a08e7422..0b8b0124e81a 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java @@ -9489,16 +9489,6 @@ List listAppExtensionAfterTimeByName( @Bind("startTime") long startTime, @Bind("extension") String extension); - @SqlQuery( - "SELECT json FROM apps_extension_time_series where appName = :appName AND extension = :extension AND timestamp >= :startTime AND timestamp < :endTime ORDER BY timestamp ASC LIMIT :limit OFFSET :offset") - List listAppExtensionInWindowByName( - @Bind("appName") String appName, - @Bind("limit") int limit, - @Bind("offset") int offset, - @Bind("startTime") long startTime, - @Bind("endTime") long endTime, - @Bind("extension") String extension); - default List listAppExtensionAfterTime( String appId, int limit, int offset, long startTime, String extension) { return listAppExtensionAfterTime(appId, limit, offset, startTime, extension, null); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfIriValidator.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfIriValidator.java new file mode 100644 index 000000000000..47ad159ae848 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfIriValidator.java @@ -0,0 +1,63 @@ +package org.openmetadata.service.rdf; + +import java.net.URI; +import java.net.URISyntaxException; +import java.util.Locale; + +/** + * Validates user-supplied IRIs before they are interpolated into SPARQL queries. + * + *

SPARQL angle-bracket-delimited IRI references must not contain characters that can escape the + * template (newlines, {@code #} comments, quotes, control characters, etc.). Stripping {@code >} + * alone is not enough — see the {@code DESCRIBE <…>} injection finding from the May 2026 PR + * review. This single utility is shared by every code path that builds a SPARQL DESCRIBE / CLEAR + * GRAPH query around a user-supplied URI so updates only need to land in one place. + * + *

Accepted form: {@code http(s)://host[/path][?query][#fragment]} with no whitespace, no angle + * brackets, no quotes, no backticks, and no control characters. Maximum 2048 characters. + */ +public final class RdfIriValidator { + + /** Soft cap. 2 KiB is well above legitimate OM URIs and keeps logs bounded. */ + static final int MAX_LENGTH = 2048; + + private RdfIriValidator() {} + + /** + * Returns the sanitized IRI when valid, {@code null} otherwise. Trims leading/trailing + * whitespace before validation; the validated form is the trimmed candidate. + */ + public static String validateEntityIri(String raw) { + if (raw == null) { + return null; + } + String candidate = raw.trim(); + if (candidate.isEmpty() || candidate.length() > MAX_LENGTH) { + return null; + } + for (int i = 0; i < candidate.length(); i++) { + char c = candidate.charAt(i); + if (c < 0x20 || c == 0x7F || c == ' ' || c == '<' || c == '>' || c == '"' || c == '\'' + || c == '`') { + return null; + } + } + try { + URI uri = new URI(candidate); + if (!uri.isAbsolute()) { + return null; + } + String scheme = uri.getScheme(); + if (scheme == null) { + return null; + } + String schemeLower = scheme.toLowerCase(Locale.ROOT); + if (!"http".equals(schemeLower) && !"https".equals(schemeLower)) { + return null; + } + } catch (URISyntaxException e) { + return null; + } + return candidate; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfUtils.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfUtils.java index e3dccc64722e..31a132987cc1 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfUtils.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/RdfUtils.java @@ -1,11 +1,13 @@ package org.openmetadata.service.rdf; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; import java.util.Set; /** * Utility methods for RDF operations */ -public class RdfUtils { +public final class RdfUtils { private static final Set PROV_ACTIVITY_TYPES = Set.of( @@ -44,9 +46,7 @@ public class RdfUtils { "dataproduct", "domain"); - private RdfUtils() { - // Private constructor for utility class - } + private RdfUtils() {} /** * Maps an entity type to its PROV-O class (Entity, Activity, or Agent). @@ -98,7 +98,30 @@ public static String getRdfType(String entityType) { case "policy" -> "om:Policy"; case "dataproduct" -> "dprod:DataProduct"; // W3C Data Product vocabulary case "domain" -> "skos:Collection"; // Organizational grouping + case "persona" -> "om:Persona"; + case "llmmodel" -> "om:LLMModel"; + case "aiapplication" -> "om:AIApplication"; + case "mcpserver" -> "om:McpServer"; + case "agentexecution" -> "om:AgentExecution"; + case "mcpexecution" -> "om:McpExecution"; + case "prompttemplate" -> "om:PromptTemplate"; + case "workflow", "workflowdefinition" -> "om:Workflow"; + case "workflowinstance" -> "om:WorkflowInstance"; + case "automation" -> "om:Automation"; default -> "om:" + entityType.substring(0, 1).toUpperCase() + entityType.substring(1); }; } + + /** + * Mints a stable, FQN-derived URI for a Column resource. Columns are sub-objects of a Table and + * have no UUID, so the FQN is the only universal identifier. The same scheme is used by the + * Table-side mapping (Table om:hasColumn) and by column-level lineage (om:fromColumn / + * om:toColumn) so that SPARQL traversal across both sides resolves to the same resource. + */ + public static String columnUri(String baseUri, String columnFqn) { + if (columnFqn == null || columnFqn.isEmpty()) { + return null; + } + return baseUri + "entity/column/" + URLEncoder.encode(columnFqn, StandardCharsets.UTF_8); + } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/extension/CustomOntologyRegistry.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/extension/CustomOntologyRegistry.java new file mode 100644 index 000000000000..35eed22b4e9b --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/extension/CustomOntologyRegistry.java @@ -0,0 +1,60 @@ +package org.openmetadata.service.rdf.extension; + +import java.util.List; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.api.configuration.rdf.CustomOntology; + +/** + * In-memory registry of user-authored ontology extensions. Each extension is keyed by its + * {@code name}. Reads are lock-free; writes are synchronized. + * + *

Persistence is intentionally deferred — admin writes that pass {@link + * CustomOntologyValidator#validate(CustomOntology)} are upserted into this registry, and the + * registry is rebuilt on server restart from any DB-backed store added in a future phase. + */ +@Slf4j +public final class CustomOntologyRegistry { + + private static final CustomOntologyRegistry INSTANCE = new CustomOntologyRegistry(); + + public static CustomOntologyRegistry getInstance() { + return INSTANCE; + } + + // ConcurrentHashMap supports lock-free reads while writes mutate concurrently. Iteration order + // is not preserved; {@link #list()} returns in whatever order ConcurrentHashMap chooses, which + // is acceptable since the only stable contract here is "all current extensions". + private final ConcurrentMap extensions = new ConcurrentHashMap<>(); + + private CustomOntologyRegistry() {} + + /** @return all extensions; iteration order is not guaranteed. */ + public List list() { + return List.copyOf(extensions.values()); + } + + public Optional get(String name) { + return Optional.ofNullable(extensions.get(name)); + } + + /** + * Insert or replace an extension. The caller is responsible for validation; this method does + * none. Returns the previous extension at that name (if any). + */ + public synchronized Optional upsert(CustomOntology extension) { + return Optional.ofNullable(extensions.put(extension.getName(), extension)); + } + + /** @return true if the extension was removed. */ + public synchronized boolean delete(String name) { + return extensions.remove(name) != null; + } + + /** Visible for tests. */ + synchronized void resetForTests() { + extensions.clear(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/extension/CustomOntologyValidator.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/extension/CustomOntologyValidator.java new file mode 100644 index 000000000000..f14d88cfb65d --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/extension/CustomOntologyValidator.java @@ -0,0 +1,283 @@ +package org.openmetadata.service.rdf.extension; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.api.configuration.rdf.CustomOntology; +import org.openmetadata.schema.api.configuration.rdf.CustomOntologyClass; +import org.openmetadata.schema.api.configuration.rdf.CustomOntologyProperty; + +/** + * Validates a {@link CustomOntology} extension before it is registered with the server. + * + *

Hard rules — the validator must reject any of these: + * + *

    + *
  1. The extension has a non-blank name and at least one class or property. + *
  2. Every custom class / property URI is in the {@code om-extension:} namespace. The + * canonical {@code om:} namespace is read-only — extensions cannot redefine + * {@code om:Column}, {@code om:Table}, etc. + *
  3. No two classes (or two properties) share the same URI within the same extension. + *
  4. Each class declares at least one parent in {@code subClassOf}; the parent must reference + * a known canonical OpenMetadata class or another class declared in this same extension. + *
  5. The class hierarchy contains no cycles ({@code A → B → A}). + *
  6. Object/Datatype property domain/range URIs reference either a canonical class or a class + * in this extension. + *
+ */ +@Slf4j +public final class CustomOntologyValidator { + + private static final String EXTENSION_NS = "https://open-metadata.org/ontology-extension/"; + private static final String CANONICAL_NS = "https://open-metadata.org/ontology/"; + + /** + * The canonical class URIs that admins are allowed to reference as parents / domains / ranges. + * Pulled from {@code openmetadata-spec/src/main/resources/rdf/ontology/openmetadata.ttl}; the + * list is intentionally small — anything outside it must be a class declared in the same + * extension. + */ + private static final Set KNOWN_CANONICAL_CLASSES = + Set.of( + "Entity", + "DataAsset", + "Service", + "Database", + "DatabaseSchema", + "Table", + "Column", + "TableConstraint", + "Pipeline", + "Topic", + "Dashboard", + "Chart", + "MLModel", + "Container", + "SearchIndex", + "APICollection", + "APIEndpoint", + "Glossary", + "GlossaryTerm", + "Tag", + "Classification", + "Domain", + "DataProduct", + "DataContract", + "Persona", + "User", + "Team", + "LineageDetails", + "ColumnLineage", + "LLMModel", + "AIApplication", + "McpServer", + "AgentExecution", + "PromptTemplate", + "Workflow", + "Automation"); + + private CustomOntologyValidator() {} + + /** @return list of validation errors. Empty list means the extension is valid. */ + public static List validate(CustomOntology extension) { + List errors = new ArrayList<>(); + if (extension == null) { + errors.add("extension must not be null"); + return errors; + } + if (isBlank(extension.getName())) { + errors.add("'name' must not be blank"); + } else if (!extension.getName().matches("^[a-z][a-z0-9-]{1,62}[a-z0-9]$")) { + errors.add( + "'name' must be 3-64 chars, lowercase letters / digits / hyphen, start with a letter"); + } + List classes = + extension.getClasses() == null ? List.of() : extension.getClasses(); + List properties = + extension.getProperties() == null ? List.of() : extension.getProperties(); + if (classes.isEmpty() && properties.isEmpty()) { + errors.add("extension must declare at least one class or property"); + } + + Set classUris = new HashSet<>(); + for (CustomOntologyClass cls : classes) { + validateClass(cls, classUris, errors); + } + + Set propertyUris = new HashSet<>(); + for (CustomOntologyProperty prop : properties) { + validateProperty(prop, propertyUris, classUris, errors); + } + + detectClassHierarchyCycles(classes, errors); + + return errors; + } + + /** @return true if validation passed; false otherwise (errors are logged at WARN). */ + public static boolean isValid(CustomOntology extension) { + List errors = validate(extension); + if (!errors.isEmpty()) { + LOG.warn( + "Custom ontology '{}' failed validation: {}", + extension == null ? "" : extension.getName(), + errors); + } + return errors.isEmpty(); + } + + private static void validateClass( + CustomOntologyClass cls, Set seenUris, List errors) { + if (cls == null) { + errors.add("null class entry"); + return; + } + if (cls.getUri() == null || cls.getUri().isBlank()) { + errors.add("class missing 'uri'"); + return; + } + if (!isExtensionUri(cls.getUri())) { + errors.add( + "class URI '" + + cls.getUri() + + "' must be in the om-extension namespace (" + + EXTENSION_NS + + "); the canonical om: namespace is read-only"); + } + if (!seenUris.add(cls.getUri())) { + errors.add("duplicate class URI in this extension: " + cls.getUri()); + } + if (cls.getSubClassOf() == null || cls.getSubClassOf().isEmpty()) { + errors.add("class '" + cls.getUri() + "' must declare at least one subClassOf parent"); + } else { + for (String parent : cls.getSubClassOf()) { + if (!isKnownClassReference(parent, seenUris)) { + errors.add( + "class '" + + cls.getUri() + + "' references unknown parent class '" + + parent + + "'; expected canonical om: class or another class in this extension"); + } + } + } + } + + private static void validateProperty( + CustomOntologyProperty prop, + Set seenUris, + Set declaredClassUris, + List errors) { + if (prop == null) { + errors.add("null property entry"); + return; + } + if (prop.getUri() == null || prop.getUri().isBlank()) { + errors.add("property missing 'uri'"); + return; + } + if (!isExtensionUri(prop.getUri())) { + errors.add("property URI '" + prop.getUri() + "' must be in the om-extension namespace"); + } + if (!seenUris.add(prop.getUri())) { + errors.add("duplicate property URI in this extension: " + prop.getUri()); + } + if (prop.getDomain() == null || prop.getDomain().isBlank()) { + errors.add("property '" + prop.getUri() + "' missing 'domain'"); + } else if (!isKnownClassReference(prop.getDomain(), declaredClassUris)) { + errors.add( + "property '" + + prop.getUri() + + "' has domain '" + + prop.getDomain() + + "' which is not a known canonical class or a class in this extension"); + } + if (prop.getRange() == null || prop.getRange().isBlank()) { + errors.add("property '" + prop.getUri() + "' missing 'range'"); + } else if (prop.getType() == CustomOntologyProperty.Type.OBJECT_PROPERTY) { + if (!isKnownClassReference(prop.getRange(), declaredClassUris)) { + errors.add( + "ObjectProperty '" + + prop.getUri() + + "' has range '" + + prop.getRange() + + "' which is not a known canonical class or a class in this extension"); + } + } else if (prop.getType() == CustomOntologyProperty.Type.DATATYPE_PROPERTY) { + if (!prop.getRange().startsWith("http://www.w3.org/2001/XMLSchema#")) { + errors.add( + "DatatypeProperty '" + + prop.getUri() + + "' range must be an xsd: datatype URI (got '" + + prop.getRange() + + "')"); + } + } + } + + /** + * Detect a cycle in the subClassOf graph using iterative DFS. A cycle is any path that returns + * to a node already on the current path stack. + */ + private static void detectClassHierarchyCycles( + List classes, List errors) { + Map> graph = new HashMap<>(); + for (CustomOntologyClass cls : classes) { + if (cls != null && cls.getUri() != null) { + graph.put(cls.getUri(), cls.getSubClassOf() == null ? List.of() : cls.getSubClassOf()); + } + } + Set visited = new HashSet<>(); + Set onStack = new HashSet<>(); + for (String node : graph.keySet()) { + if (hasCycle(node, graph, visited, onStack)) { + errors.add("class hierarchy contains a cycle through " + node); + } + } + } + + private static boolean hasCycle( + String node, Map> graph, Set visited, Set onStack) { + if (onStack.contains(node)) return true; + if (visited.contains(node)) return false; + onStack.add(node); + visited.add(node); + for (String parent : graph.getOrDefault(node, List.of())) { + if (graph.containsKey(parent) && hasCycle(parent, graph, visited, onStack)) { + return true; + } + } + onStack.remove(node); + return false; + } + + private static boolean isExtensionUri(String uri) { + return uri != null && uri.startsWith(EXTENSION_NS); + } + + /** + * @return true if the URI references either a canonical om: class on the allowlist or a class + * declared inside the current extension. + */ + private static boolean isKnownClassReference(String uri, Set extensionClassUris) { + if (uri == null) return false; + if (extensionClassUris.contains(uri)) return true; + if (uri.startsWith(CANONICAL_NS)) { + String localName = uri.substring(CANONICAL_NS.length()); + return KNOWN_CANONICAL_CLASSES.contains(localName); + } + // Allow short-form prefixed names like "om:Table". + if (uri.startsWith("om:")) { + return KNOWN_CANONICAL_CLASSES.contains(uri.substring(3)); + } + return false; + } + + private static boolean isBlank(String s) { + return s == null || s.isBlank(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/federation/SparqlFederationGuard.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/federation/SparqlFederationGuard.java new file mode 100644 index 000000000000..0fa99160e48b --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/federation/SparqlFederationGuard.java @@ -0,0 +1,187 @@ +package org.openmetadata.service.rdf.federation; + +import java.net.URI; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; +import lombok.extern.slf4j.Slf4j; +import org.apache.jena.query.Query; +import org.apache.jena.query.QueryException; +import org.apache.jena.query.QueryFactory; +import org.apache.jena.sparql.syntax.ElementService; +import org.apache.jena.sparql.syntax.ElementSubQuery; +import org.apache.jena.sparql.syntax.ElementVisitorBase; +import org.apache.jena.sparql.syntax.ElementWalker; +import org.openmetadata.schema.api.configuration.rdf.RdfConfiguration; +import org.openmetadata.schema.api.configuration.rdf.SparqlFederationConfig; + +/** + * Inspects an incoming SPARQL query for {@code SERVICE } clauses and rejects any whose + * endpoint URI is not in the configured allowlist. + * + *

Detection uses Jena's {@link ElementWalker} so it sees only real SERVICE elements — the + * keyword "SERVICE" inside a string literal or a comment is correctly ignored. The walker also + * recurses into subqueries, OPTIONAL/UNION/MINUS branches, and nested SERVICE blocks. + * + *

Behavior: + * + *

    + *
  • Federation disabled (the default) — any SERVICE clause is a violation. + *
  • Federation enabled — a SERVICE clause must reference a URI present in + * {@code allowedEndpoints} verbatim. + *
  • Variable SERVICE endpoints ({@code SERVICE ?endpoint}) cannot be statically allowlisted + * and are always rejected. + *
  • Queries that fail to parse here are passed through; the SPARQL engine returns its own + * parse error to the caller, preserving message fidelity. + *
+ */ +@Slf4j +public final class SparqlFederationGuard { + + private final boolean federationEnabled; + private final Set allowedEndpoints; + + public SparqlFederationGuard(RdfConfiguration config) { + SparqlFederationConfig federation = config == null ? null : config.getFederation(); + this.federationEnabled = federation != null && Boolean.TRUE.equals(federation.getEnabled()); + this.allowedEndpoints = + federation == null || federation.getAllowedEndpoints() == null + ? Set.of() + : federation.getAllowedEndpoints().stream() + .map(URI::toString) + .collect(Collectors.toUnmodifiableSet()); + } + + /** Visible for tests. Package-private constructor that takes the policy directly. */ + SparqlFederationGuard(boolean federationEnabled, Set allowedEndpoints) { + this.federationEnabled = federationEnabled; + this.allowedEndpoints = allowedEndpoints == null ? Set.of() : Set.copyOf(allowedEndpoints); + } + + /** + * @return all distinct SERVICE endpoint URIs found in the query. Order of first appearance is + * preserved. Variable endpoints surface as the literal string {@code ?varname}. + */ + public List serviceEndpoints(String sparql) { + Optional parsed = parseQuietly(sparql); + if (parsed.isEmpty()) { + return List.of(); + } + EndpointCollector collector = new EndpointCollector(); + ElementWalker.walk(parsed.get().getQueryPattern(), collector); + return List.copyOf(collector.endpoints); + } + + /** + * @return the first endpoint that violates the policy, or empty if the query is allowed. + */ + public Optional firstDisallowedEndpoint(String sparql) { + for (String endpoint : serviceEndpoints(sparql)) { + if (!isAllowed(endpoint)) { + return Optional.of(endpoint); + } + } + return Optional.empty(); + } + + /** + * Convenience: throw {@link FederationDisallowedException} if any SERVICE clause is rejected. + */ + public void enforce(String sparql) { + Optional blocked = firstDisallowedEndpoint(sparql); + if (blocked.isPresent()) { + throw new FederationDisallowedException(blocked.get(), federationEnabled, allowedEndpoints); + } + } + + private boolean isAllowed(String endpoint) { + if (endpoint.startsWith("?")) { + // Variable endpoints can't be statically allowlisted. + return false; + } + if (!federationEnabled) { + return false; + } + return allowedEndpoints.contains(endpoint); + } + + private Optional parseQuietly(String sparql) { + try { + return Optional.ofNullable(QueryFactory.create(sparql)); + } catch (QueryException e) { + LOG.debug( + "SPARQL parse failed inside federation guard; deferring to engine: {}", e.getMessage()); + return Optional.empty(); + } + } + + private static final class EndpointCollector extends ElementVisitorBase { + private final Set endpoints = new LinkedHashSet<>(); + + @Override + public void visit(ElementService el) { + if (el.getServiceNode().isVariable()) { + endpoints.add("?" + el.getServiceNode().getName()); + } else if (el.getServiceNode().isURI()) { + endpoints.add(el.getServiceNode().getURI()); + } + } + + @Override + public void visit(ElementSubQuery el) { + // ElementWalker stops at subquery boundaries; descend manually so a SERVICE inside an + // inner SELECT still gets caught. + if (el.getQuery() != null && el.getQuery().getQueryPattern() != null) { + ElementWalker.walk(el.getQuery().getQueryPattern(), this); + } + } + } + + /** + * Thrown by {@link #enforce(String)} when a query references a disallowed endpoint. Carries the + * effective policy so callers can include it in the error response. + */ + public static final class FederationDisallowedException extends RuntimeException { + + private final String blockedEndpoint; + private final boolean federationEnabled; + private final Set allowedEndpoints; + + FederationDisallowedException( + String blockedEndpoint, boolean federationEnabled, Set allowedEndpoints) { + super(buildMessage(blockedEndpoint, federationEnabled, allowedEndpoints)); + this.blockedEndpoint = blockedEndpoint; + this.federationEnabled = federationEnabled; + this.allowedEndpoints = Collections.unmodifiableSet(allowedEndpoints); + } + + public String getBlockedEndpoint() { + return blockedEndpoint; + } + + public boolean isFederationEnabled() { + return federationEnabled; + } + + public Set getAllowedEndpoints() { + return allowedEndpoints; + } + + private static String buildMessage( + String endpoint, boolean federationEnabled, Set allowedEndpoints) { + if (!federationEnabled) { + return "SPARQL SERVICE clause references " + + endpoint + + " but federated SPARQL is disabled. " + + "Enable rdf.federation.enabled and add the endpoint to rdf.federation.allowedEndpoints."; + } + return "SPARQL SERVICE clause references " + + endpoint + + " which is not in the allowlist. Allowed endpoints: " + + allowedEndpoints; + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/inference/InferenceRuleRegistry.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/inference/InferenceRuleRegistry.java new file mode 100644 index 000000000000..3b6f6661887f --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/inference/InferenceRuleRegistry.java @@ -0,0 +1,114 @@ +package org.openmetadata.service.rdf.inference; + +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.IOException; +import java.io.InputStream; +import java.util.Comparator; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.api.configuration.rdf.InferenceRule; + +/** + * In-memory registry of OpenMetadata inference rules. The starter pack is loaded once per JVM + * from the classpath under {@code rdf/inference-rules/}; further entries can be {@link + * #upsert(InferenceRule) upserted} programmatically. + * + *

This is intentionally a Phase-1 storage shape — there is no DB-backed persistence yet. + * Admin REST writes call {@link #upsert(InferenceRule)} after passing validation. A follow-up + * phase will swap this in-memory registry for a JDBI-backed repository without changing the + * exposed API. + */ +@Slf4j +public final class InferenceRuleRegistry { + + private static final String[] STARTER_PACK = { + "/rdf/inference-rules/transitive-lineage-closure.json", + "/rdf/inference-rules/pii-propagation-via-lineage.json", + "/rdf/inference-rules/schema-tag-inheritance.json", + "/rdf/inference-rules/domain-membership-inheritance.json" + }; + + private static final InferenceRuleRegistry INSTANCE = new InferenceRuleRegistry(); + + public static InferenceRuleRegistry getInstance() { + return INSTANCE; + } + + // ConcurrentHashMap supports lock-free reads from {@link #list()} / {@link #get(String)} while + // {@link #upsert} / {@link #delete} mutate concurrently. Iteration order isn't preserved, so + // {@code list()} sorts explicitly (priority + name) for deterministic API output. + private final ConcurrentMap rules = new ConcurrentHashMap<>(); + private final ObjectMapper mapper = new ObjectMapper(); + private volatile boolean starterLoaded = false; + + private InferenceRuleRegistry() {} + + /** + * Load the starter pack from the classpath. Idempotent. + */ + public synchronized void loadStarterPackIfNeeded() { + if (starterLoaded) return; + int loaded = 0; + for (String path : STARTER_PACK) { + try (InputStream is = InferenceRuleRegistry.class.getResourceAsStream(path)) { + if (is == null) { + LOG.warn("Inference rule starter pack resource missing: {}", path); + continue; + } + InferenceRule rule = mapper.readValue(is, InferenceRule.class); + List errors = InferenceRuleValidator.validate(rule); + if (!errors.isEmpty()) { + LOG.error("Starter pack rule '{}' failed validation: {}", path, errors); + continue; + } + rules.put(rule.getName(), rule); + loaded++; + } catch (IOException e) { + LOG.error("Failed to load starter pack rule {}", path, e); + } + } + starterLoaded = true; + LOG.info("Loaded {} inference rules from starter pack", loaded); + } + + /** @return all rules in priority order, then by name. */ + public List list() { + loadStarterPackIfNeeded(); + return rules.values().stream() + .sorted( + Comparator.comparing( + (InferenceRule r) -> r.getPriority() == null ? 100 : r.getPriority()) + .thenComparing(InferenceRule::getName)) + .toList(); + } + + /** @return the rule with the given name, or empty if no such rule. */ + public Optional get(String name) { + loadStarterPackIfNeeded(); + return Optional.ofNullable(rules.get(name)); + } + + /** + * Insert or replace a rule. The caller is responsible for validation (see {@link + * InferenceRuleValidator}); this method does no validation itself. + */ + public synchronized void upsert(InferenceRule rule) { + loadStarterPackIfNeeded(); + rules.put(rule.getName(), rule); + } + + /** @return true if a rule with that name was removed. */ + public synchronized boolean delete(String name) { + loadStarterPackIfNeeded(); + return rules.remove(name) != null; + } + + /** Clear the registry. Visible for tests. */ + synchronized void resetForTests() { + rules.clear(); + starterLoaded = false; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/inference/InferenceRuleValidator.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/inference/InferenceRuleValidator.java new file mode 100644 index 000000000000..4ccb49413db4 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/inference/InferenceRuleValidator.java @@ -0,0 +1,142 @@ +package org.openmetadata.service.rdf.inference; + +import java.util.ArrayList; +import java.util.List; +import lombok.extern.slf4j.Slf4j; +import org.apache.jena.query.Query; +import org.apache.jena.query.QueryException; +import org.apache.jena.query.QueryFactory; +import org.apache.jena.sparql.syntax.ElementService; +import org.apache.jena.sparql.syntax.ElementSubQuery; +import org.apache.jena.sparql.syntax.ElementVisitorBase; +import org.apache.jena.sparql.syntax.ElementWalker; +import org.openmetadata.schema.api.configuration.rdf.InferenceRule; + +/** + * Validates {@link InferenceRule} payloads before they are accepted or executed. + * + *

The validator is intentionally strict — admins write these rules and they run server-side + * against the whole graph, so a malformed or hostile rule has wide blast radius. + * + *

Checks performed (in order; first failing check returns): + * + *

    + *
  1. The rule has a non-blank name and rule body. + *
  2. The body parses as a SPARQL Query (rejects SPARQL UPDATE — those are emitted via the + * resulting CONSTRUCT triples, not by the rule body itself). + *
  3. For {@code ruleType=CONSTRUCT}, the parsed query must be CONSTRUCT type. SELECT, ASK, + * and DESCRIBE rules are rejected — they don't produce inferable triples. + *
  4. The body must not contain any SERVICE clauses. Inference must be deterministic and run + * against the local graph; federated lookups are rejected. + *
  5. The body must not be a no-op CONSTRUCT (empty WHERE) — those would either produce + * nothing or, with ASK semantics, blow up. + *
  6. Priority is within bounds. + *
+ */ +@Slf4j +public final class InferenceRuleValidator { + + private InferenceRuleValidator() {} + + /** @return the list of validation errors. Empty list means the rule is valid. */ + public static List validate(InferenceRule rule) { + List errors = new ArrayList<>(); + if (rule == null) { + errors.add("rule must not be null"); + return errors; + } + if (isBlank(rule.getName())) { + errors.add("'name' must not be blank"); + } else if (!rule.getName().matches("^[a-z][a-z0-9-]{1,62}[a-z0-9]$")) { + errors.add( + "'name' must be 3-64 chars, lowercase letters / digits / hyphen, start with a letter, end with a letter or digit"); + } + if (isBlank(rule.getRuleBody())) { + errors.add("'ruleBody' must not be blank"); + } + InferenceRule.RuleType ruleType = + rule.getRuleType() == null ? InferenceRule.RuleType.CONSTRUCT : rule.getRuleType(); + if (ruleType == InferenceRule.RuleType.RDFS) { + // RDFS is a placeholder — the engine doesn't ship a parser for that body shape yet. + errors.add( + "ruleType=RDFS is reserved for future use; current engine only ships CONSTRUCT support"); + return errors; + } + if (rule.getPriority() != null && (rule.getPriority() < 0 || rule.getPriority() > 10_000)) { + errors.add("'priority' must be between 0 and 10000"); + } + if (errors.size() > 0) { + return errors; + } + + Query parsed; + try { + parsed = QueryFactory.create(rule.getRuleBody()); + } catch (QueryException e) { + errors.add("ruleBody failed to parse as SPARQL: " + e.getMessage()); + return errors; + } + if (!parsed.isConstructType()) { + errors.add( + "ruleBody must be a SPARQL CONSTRUCT query for ruleType=CONSTRUCT (got " + + parsed.queryType() + + "); inference rules emit new triples and only CONSTRUCT does that"); + return errors; + } + if (parsed.getQueryPattern() == null || isEmptyPattern(parsed.getQueryPattern().toString())) { + errors.add("ruleBody must have a non-empty WHERE pattern"); + } + if (parsed.getConstructTemplate() == null + || parsed.getConstructTemplate().getTriples().isEmpty()) { + errors.add("ruleBody CONSTRUCT template must contain at least one triple pattern"); + } + ServiceFinder serviceFinder = new ServiceFinder(); + if (parsed.getQueryPattern() != null) { + ElementWalker.walk(parsed.getQueryPattern(), serviceFinder); + } + if (serviceFinder.found) { + errors.add( + "ruleBody must not contain SERVICE clauses; inference is local-only and federated rules are rejected"); + } + return errors; + } + + /** + * @return true if the rule passed validation; false otherwise. Errors are logged at WARN. + */ + public static boolean isValid(InferenceRule rule) { + List errors = validate(rule); + if (!errors.isEmpty()) { + LOG.warn( + "Inference rule '{}' failed validation: {}", + rule == null ? "" : rule.getName(), + errors); + } + return errors.isEmpty(); + } + + private static boolean isBlank(String s) { + return s == null || s.isBlank(); + } + + private static boolean isEmptyPattern(String pattern) { + String trimmed = pattern.replaceAll("\\s", ""); + return trimmed.isEmpty() || trimmed.equals("{}"); + } + + private static final class ServiceFinder extends ElementVisitorBase { + boolean found; + + @Override + public void visit(ElementService el) { + found = true; + } + + @Override + public void visit(ElementSubQuery el) { + if (el.getQuery() != null && el.getQuery().getQueryPattern() != null) { + ElementWalker.walk(el.getQuery().getQueryPattern(), this); + } + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CentralityComputation.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CentralityComputation.java new file mode 100644 index 000000000000..2d7de87dd8b8 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CentralityComputation.java @@ -0,0 +1,188 @@ +package org.openmetadata.service.rdf.insights; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.rdf.RdfRepository; + +/** + * Pulls a graph snapshot of one entity type out of Fuseki, runs {@link PageRank} on it, and + * persists the resulting scores back into the named graph + * {@code }. The {@code /v1/rdf/insights/important} + * endpoint reads those triples through the {@code om:centralityScore} predicate. + * + *

Edge weights chosen to reflect importance for governance purposes: + * + *

    + *
  • {@code prov:wasDerivedFrom} — 1.0 (lineage) + *
  • {@code om:hasTag}, {@code om:hasGlossaryTerm} — 0.5 (semantic linkage) + *
  • {@code om:hasColumn} — 0.2 (containment, weak) + *
+ * + *

Scope: only entities of the requested type are added as nodes; off-type targets become + * dangling sinks so they don't hijack mass. PageRank is run with default damping 0.85. + */ +@Slf4j +public final class CentralityComputation { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String INSIGHTS_GRAPH_PREFIX = OM_NS + "insights/centrality/"; + + private final RdfRepository repository; + private final PageRank pageRank; + + public CentralityComputation(RdfRepository repository) { + this(repository, new PageRank()); + } + + CentralityComputation(RdfRepository repository, PageRank pageRank) { + this.repository = repository; + this.pageRank = pageRank; + } + + /** + * Run centrality for one entity type end-to-end: extract → compute → persist. + * + * @return summary describing what got computed. + */ + public Result computeAndPersist(String entityType) { + String safeType = ImportanceQueryBuilder.validateEntityType(entityType); + String classLocalName = capitalize(safeType); + + Map> graph = extractGraph(classLocalName); + if (graph.isEmpty()) { + LOG.info("No entities of type {} found; skipping centrality run", classLocalName); + return new Result(safeType, 0, 0, false); + } + PageRank.Result ranked = pageRank.compute(graph); + persistScores(classLocalName, ranked.scores()); + return new Result(safeType, ranked.scores().size(), ranked.iterations(), ranked.converged()); + } + + /** Extract the graph snapshot from Fuseki via SPARQL. Visible for testing. */ + Map> extractGraph(String classLocalName) { + String sparql = + String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "PREFIX prov: ", + "SELECT ?from ?to ?predicate WHERE {", + " ?from a om:" + classLocalName + " .", + " {", + " ?from prov:wasDerivedFrom ?to .", + " BIND(\"prov:wasDerivedFrom\" AS ?predicate)", + " } UNION {", + " ?from om:hasTag ?to .", + " BIND(\"om:hasTag\" AS ?predicate)", + " } UNION {", + " ?from om:hasGlossaryTerm ?to .", + " BIND(\"om:hasGlossaryTerm\" AS ?predicate)", + " } UNION {", + " ?from om:hasColumn ?to .", + " BIND(\"om:hasColumn\" AS ?predicate)", + " }", + "}"); + + String json; + try { + json = repository.executeSparqlQuery(sparql, "application/sparql-results+json"); + } catch (Exception e) { + LOG.error("Failed to extract centrality graph for {}", classLocalName, e); + return Map.of(); + } + return parseGraph(json); + } + + /** Parse SPARQL JSON results into a weighted adjacency map. Visible for testing. */ + static Map> parseGraph(String selectJson) { + Map> graph = new LinkedHashMap<>(); + if (selectJson == null || selectJson.isBlank()) { + return graph; + } + try { + JsonNode root = JsonUtils.readTree(selectJson); + JsonNode bindings = root.path("results").path("bindings"); + if (!bindings.isArray()) return graph; + for (JsonNode row : bindings) { + String from = textValue(row, "from"); + String to = textValue(row, "to"); + String predicate = textValue(row, "predicate"); + if (from == null || to == null || predicate == null) continue; + double weight = weightFor(predicate); + graph.computeIfAbsent(from, k -> new HashMap<>()).merge(to, weight, Double::sum); + } + } catch (Exception e) { + LOG.warn("Failed to parse centrality SPARQL result: {}", e.getMessage()); + } + return graph; + } + + static double weightFor(String predicate) { + return switch (predicate) { + case "prov:wasDerivedFrom" -> 1.0; + case "om:hasTag", "om:hasGlossaryTerm" -> 0.5; + case "om:hasColumn" -> 0.2; + default -> 0.0; + }; + } + + /** Write the scores back to Fuseki under a dedicated named graph. Visible for testing. */ + void persistScores(String classLocalName, Map scores) { + String graphUri = INSIGHTS_GRAPH_PREFIX + classLocalName.toLowerCase(java.util.Locale.ROOT); + StringBuilder update = new StringBuilder(); + update.append("PREFIX om: <").append(OM_NS).append(">\n"); + update.append("PREFIX xsd: \n"); + update + .append("WITH <") + .append(graphUri) + .append( + "> DELETE { ?s om:centralityScore ?o ; om:centralityRank ?r } WHERE { ?s om:centralityScore ?o . OPTIONAL { ?s om:centralityRank ?r } } ;\n"); + update.append("INSERT DATA { GRAPH <").append(graphUri).append("> {\n"); + int rank = 1; + for (Map.Entry e : + scores.entrySet().stream() + .sorted(Map.Entry.comparingByValue().reversed()) + .toList()) { + update + .append(" <") + .append(e.getKey()) + .append("> om:centralityScore \"") + .append(e.getValue()) + .append("\"^^xsd:double ; om:centralityRank \"") + .append(rank++) + .append("\"^^xsd:integer .\n"); + } + update.append("} }"); + try { + repository.executeSparqlUpdate(update.toString()); + } catch (Exception e) { + LOG.error("Failed to persist centrality scores for {}", classLocalName, e); + } + } + + private static String textValue(JsonNode row, String varName) { + JsonNode node = row.path(varName); + if (node.isMissingNode() || node.isNull()) return null; + JsonNode value = node.path("value"); + return value.isMissingNode() || value.isNull() ? null : value.asText(); + } + + private static String capitalize(String s) { + if (s == null || s.isEmpty()) return s; + return Character.toUpperCase(s.charAt(0)) + s.substring(1); + } + + /** Result of a centrality run. */ + public record Result(String entityType, int nodesScored, int iterations, boolean converged) {} + + /** Helper for tests to construct rows for parseGraph. */ + static List exampleRows() { + return List.of( + new String[] {"urn:t1", "urn:t2", "prov:wasDerivedFrom"}, + new String[] {"urn:t1", "urn:c1", "om:hasColumn"}); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CoOccurrenceQueryBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CoOccurrenceQueryBuilder.java new file mode 100644 index 000000000000..aed4cbe85d2b --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CoOccurrenceQueryBuilder.java @@ -0,0 +1,101 @@ +package org.openmetadata.service.rdf.insights; + +/** + * SPARQL builders for Phase 3.5 catalog-wide insights — pure aggregate views over the existing + * triples, no precomputation. Each method is a standalone SPARQL SELECT with input validation; + * callers ({@code RdfResource}) hand the query verbatim to the SPARQL endpoint and stream the + * results back as SPARQL-JSON. + * + *

    + *
  • {@link #tagCoOccurrence(int, int)} — pairs of tags that get applied to the same entity, by + * overlap count. Pairs are canonicalised (str(a) < str(b)) so each pair is reported once. + *
  • {@link #glossaryReach(int, int)} — glossary terms ranked by the number of distinct domains + * they appear under. Useful for "term used across the most domains" insight in Phase 3.5. + *
+ */ +public final class CoOccurrenceQueryBuilder { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + + /** Default top-N when caller doesn't specify. */ + public static final int DEFAULT_LIMIT = 20; + + /** Hard cap so a buggy caller can't ask for tens of thousands of rows. */ + public static final int MAX_LIMIT = 100; + + /** Default minimum overlap threshold — pairs that co-occur on fewer entities are dropped. */ + public static final int DEFAULT_MIN_COUNT = 2; + + private CoOccurrenceQueryBuilder() {} + + /** + * Tag co-occurrence: pairs of tags applied to the same entity together at least {@code minCount} + * times. Result columns: ?tagA, ?tagB, ?count. + * + * @param minCount minimum number of shared entities; values < 1 are clamped to 1 + * @param limit number of rows; clamped to [1, {@link #MAX_LIMIT}] + */ + public static String tagCoOccurrence(int minCount, int limit) { + int safeMin = clamp(minCount, 1, Integer.MAX_VALUE); + int safeLimit = clamp(limit, 1, MAX_LIMIT); + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "SELECT ?tagA ?tagB (COUNT(?entity) AS ?count) WHERE {", + " ?entity om:hasTag ?tagA .", + " ?entity om:hasTag ?tagB .", + " FILTER(STR(?tagA) < STR(?tagB))", + "}", + "GROUP BY ?tagA ?tagB", + "HAVING (COUNT(?entity) >= " + safeMin + ")", + "ORDER BY DESC(?count) ?tagA ?tagB", + "LIMIT " + safeLimit); + } + + /** + * Glossary reach: each glossary term + the number of distinct domains it shows up in. A term + * that's used by tables across many domains is more cross-cutting and signals a richer concept. + * + *

Result columns: ?term, ?domainCount. + * + * @param minDomains floor on the count; values < 1 become 1 + * @param limit number of rows; clamped to [1, {@link #MAX_LIMIT}] + */ + public static String glossaryReach(int minDomains, int limit) { + int safeMin = clamp(minDomains, 1, Integer.MAX_VALUE); + int safeLimit = clamp(limit, 1, MAX_LIMIT); + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "SELECT ?term (COUNT(DISTINCT ?domain) AS ?domainCount) WHERE {", + " ?entity om:hasGlossaryTerm ?term .", + " ?entity om:hasDomain ?domain .", + "}", + "GROUP BY ?term", + "HAVING (COUNT(DISTINCT ?domain) >= " + safeMin + ")", + "ORDER BY DESC(?domainCount) ?term", + "LIMIT " + safeLimit); + } + + /** + * Tag popularity: tags ranked by the number of entities they're applied to. Result columns: + * ?tag, ?entityCount. + */ + public static String tagPopularity(int limit) { + int safeLimit = clamp(limit, 1, MAX_LIMIT); + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "SELECT ?tag (COUNT(DISTINCT ?entity) AS ?entityCount) WHERE {", + " ?entity om:hasTag ?tag .", + "}", + "GROUP BY ?tag", + "ORDER BY DESC(?entityCount) ?tag", + "LIMIT " + safeLimit); + } + + private static int clamp(int v, int lo, int hi) { + if (v < lo) return lo; + return Math.min(v, hi); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CommunityComputation.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CommunityComputation.java new file mode 100644 index 000000000000..d4946fac4a33 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/CommunityComputation.java @@ -0,0 +1,281 @@ +package org.openmetadata.service.rdf.insights; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.rdf.RdfRepository; + +/** + * Phase 3.2 community-detection driver. Pulls a sub-graph out of Fuseki for one + * {@link GraphType graph kind} (lineage or tag-co-occurrence) restricted to one entity type, runs + * {@link Louvain}, and persists the resulting partition under a dedicated named graph + * {@code }. + * + *

Each persisted community is an {@code om:Community} resource with {@code om:hasMember}, + * {@code om:modularity}, {@code om:communityType}, and {@code om:communitySize}. The + * {@code modularity} value is the same on every community of a given run — it describes the + * partition as a whole, not any single cluster. + * + *

Determinism: {@code Louvain} processes nodes in the SPARQL result iteration order, so as long + * as Fuseki returns rows in a stable order (Jena's TDB does, modulo equal-cost solutions), the + * persisted membership is reproducible. + */ +@Slf4j +public final class CommunityComputation { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String INSIGHTS_GRAPH_PREFIX = OM_NS + "insights/communities/"; + private static final String COMMUNITY_URI_PREFIX = OM_NS + "instance/Community/"; + + private final RdfRepository repository; + private final Louvain louvain; + + public CommunityComputation(RdfRepository repository) { + this(repository, new Louvain()); + } + + CommunityComputation(RdfRepository repository, Louvain louvain) { + this.repository = repository; + this.louvain = louvain; + } + + public Result computeAndPersist(String entityType, String graphType) { + String safeType = ImportanceQueryBuilder.validateEntityType(entityType); + GraphType gt = GraphType.parse(graphType); + String classLocalName = capitalize(safeType); + + Map> graph = extractGraph(classLocalName, gt); + if (graph.isEmpty()) { + LOG.info( + "No edges of kind {} found for entity type {}; skipping community run", + gt, + classLocalName); + return new Result(safeType, gt.label, 0, 0, 0.0); + } + Louvain.Result partition = louvain.compute(graph); + persistCommunities(classLocalName, gt, partition); + return new Result( + safeType, + gt.label, + partition.communityCount(), + partition.communityByNode().size(), + partition.modularity()); + } + + /** + * Build the SPARQL SELECT used by GET /v1/rdf/insights/communities to list previously persisted + * communities for the given (entityType, graphType) pair. + */ + public static String listingSparql(String entityType, String graphType) { + String safeType = ImportanceQueryBuilder.validateEntityType(entityType); + GraphType gt = GraphType.parse(graphType); + String classLocalName = capitalize(safeType); + String graphUri = + INSIGHTS_GRAPH_PREFIX + gt.label + "/" + classLocalName.toLowerCase(Locale.ROOT); + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "SELECT ?community ?size ?modularity ?member", + "FROM <" + graphUri + ">", + "WHERE {", + " ?community a om:Community ;", + " om:communitySize ?size ;", + " om:modularity ?modularity ;", + " om:hasMember ?member .", + "}", + "ORDER BY DESC(?size) ?community ?member"); + } + + Map> extractGraph(String classLocalName, GraphType graphType) { + String sparql = + switch (graphType) { + case LINEAGE -> lineageGraphSparql(classLocalName); + case TAG_CO_OCCURRENCE -> tagCoOccurrenceSparql(classLocalName); + }; + String json; + try { + json = repository.executeSparqlQuery(sparql, "application/sparql-results+json"); + } catch (Exception e) { + LOG.error("Failed to extract {} graph for {}", graphType, classLocalName, e); + return Map.of(); + } + return parseGraph(json); + } + + static String lineageGraphSparql(String classLocalName) { + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "PREFIX prov: ", + "SELECT ?from ?to (1.0 AS ?weight) WHERE {", + " ?from a om:" + classLocalName + " .", + " ?to a om:" + classLocalName + " .", + " {", + " ?from prov:wasDerivedFrom ?to .", + " } UNION {", + " ?from om:upstream ?to .", + " } UNION {", + " ?to om:downstream ?from .", + " }", + " FILTER(?from != ?to)", + "}"); + } + + static String tagCoOccurrenceSparql(String classLocalName) { + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "SELECT ?from ?to (COUNT(?shared) AS ?weight) WHERE {", + " ?from a om:" + classLocalName + " .", + " ?to a om:" + classLocalName + " .", + " {", + " ?from om:hasTag ?shared .", + " ?to om:hasTag ?shared .", + " } UNION {", + " ?from om:hasGlossaryTerm ?shared .", + " ?to om:hasGlossaryTerm ?shared .", + " }", + " FILTER(STR(?from) < STR(?to))", + "}", + "GROUP BY ?from ?to"); + } + + /** + * Parses the SPARQL bindings into a directed adjacency map (one weighted entry per edge). The + * companion SPARQL query already canonicalises with {@code FILTER(STR(?from) < STR(?to))} so + * each pair appears once. {@link Louvain#addAllEdges} is responsible for symmetrising the + * adjacency internally — emitting both directions here would double-count every edge weight. + * + *

Each target node is still registered as a (possibly-empty) key so that Louvain's + * {@code graph.keySet()} includes every participating node, not only the ones that appear as a + * source. + */ + static Map> parseGraph(String selectJson) { + Map> graph = new LinkedHashMap<>(); + if (selectJson == null || selectJson.isBlank()) return graph; + try { + JsonNode root = JsonUtils.readTree(selectJson); + JsonNode bindings = root.path("results").path("bindings"); + if (!bindings.isArray()) return graph; + for (JsonNode row : bindings) { + String from = textValue(row, "from"); + String to = textValue(row, "to"); + if (from == null || to == null || from.equals(to)) continue; + double weight = doubleValue(row, "weight", 1.0); + if (weight <= 0) continue; + graph.computeIfAbsent(from, k -> new HashMap<>()).merge(to, weight, Double::sum); + // Register the target as a node too, without adding the reverse edge weight. Louvain + // will symmetrise the adjacency itself; we just need every node visible to it. + graph.computeIfAbsent(to, k -> new HashMap<>()); + } + } catch (Exception e) { + LOG.warn("Failed to parse community graph SPARQL result: {}", e.getMessage()); + } + return graph; + } + + void persistCommunities( + String classLocalName, GraphType graphType, Louvain.Result partition) { + String graphUri = + INSIGHTS_GRAPH_PREFIX + graphType.label + "/" + classLocalName.toLowerCase(Locale.ROOT); + Map> members = partition.membersByCommunity(); + + StringBuilder update = new StringBuilder(); + update.append("PREFIX om: <").append(OM_NS).append(">\n"); + update.append("PREFIX xsd: \n"); + update + .append("WITH <") + .append(graphUri) + .append("> DELETE { ?s ?p ?o } WHERE { ?s a om:Community ; ?p ?o } ;\n"); + update.append("INSERT DATA { GRAPH <").append(graphUri).append("> {\n"); + for (Map.Entry> entry : members.entrySet()) { + String communityUri = + COMMUNITY_URI_PREFIX + + graphType.label + + "/" + + classLocalName.toLowerCase(Locale.ROOT) + + "/" + + entry.getKey(); + update + .append(" <") + .append(communityUri) + .append("> a om:Community ; om:communityType \"") + .append(graphType.label) + .append("\" ; om:communitySize \"") + .append(entry.getValue().size()) + .append("\"^^xsd:integer ; om:modularity \"") + .append(partition.modularity()) + .append("\"^^xsd:double"); + for (String member : entry.getValue()) { + update.append(" ; om:hasMember <").append(member).append(">"); + } + update.append(" .\n"); + } + update.append("} }"); + try { + repository.executeSparqlUpdate(update.toString()); + } catch (Exception e) { + LOG.error("Failed to persist communities for {} / {}", classLocalName, graphType, e); + } + } + + private static String textValue(JsonNode row, String varName) { + JsonNode node = row.path(varName); + if (node.isMissingNode() || node.isNull()) return null; + JsonNode value = node.path("value"); + return value.isMissingNode() || value.isNull() ? null : value.asText(); + } + + private static double doubleValue(JsonNode row, String varName, double fallback) { + JsonNode node = row.path(varName); + if (node.isMissingNode() || node.isNull()) return fallback; + JsonNode value = node.path("value"); + if (value.isMissingNode() || value.isNull()) return fallback; + try { + return Double.parseDouble(value.asText()); + } catch (NumberFormatException e) { + return fallback; + } + } + + private static String capitalize(String s) { + if (s == null || s.isEmpty()) return s; + return Character.toUpperCase(s.charAt(0)) + s.substring(1); + } + + /** Source graph kinds supported by community detection. */ + public enum GraphType { + LINEAGE("lineage"), + TAG_CO_OCCURRENCE("tagCoOccurrence"); + + public final String label; + + GraphType(String label) { + this.label = label; + } + + public static GraphType parse(String value) { + if (value == null || value.isBlank()) return LINEAGE; + String norm = value.trim().toLowerCase(Locale.ROOT); + return switch (norm) { + case "lineage" -> LINEAGE; + case "tagcooccurrence", + "tag", + "tags", + "tag-co-occurrence", + "tag_co_occurrence" -> TAG_CO_OCCURRENCE; + default -> throw new IllegalArgumentException( + "graphType must be one of: lineage, tagCoOccurrence (got: " + value + ")"); + }; + } + } + + /** Result of a community-detection run. */ + public record Result( + String entityType, String graphType, int communities, int membersTotal, double modularity) {} +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/ImportanceQueryBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/ImportanceQueryBuilder.java new file mode 100644 index 000000000000..3c71a529a9c4 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/ImportanceQueryBuilder.java @@ -0,0 +1,111 @@ +package org.openmetadata.service.rdf.insights; + +/** + * Builds the SPARQL query that ranks entities by an importance score blending OpenMetadata's + * existing usage percentile (real query data) with downstream lineage count (graph topology). + * + *

Scoring formula: + * + * 

{@code
+ * score = 0.6 * (usagePercentile / 100)
+ *       + 0.4 * (downstreamCount / max(downstreamCount across the entity type))
+ *       + 0.0 * centralityScore   // 3.1.b will plug PageRank here for null-usage entities
+ * }
+ * + *

Both terms are 0–1 after normalization. Entities without usage data fall to the bottom + * until 3.1.b's PageRank fallback lands and {@code om:centralityScore} starts populating. + */ +public final class ImportanceQueryBuilder { + + /** Hard-capped to keep the response page-sized. */ + static final int MIN_LIMIT = 1; + + static final int MAX_LIMIT = 100; + static final int DEFAULT_LIMIT = 20; + + /** Allowed window values map to the matching `om:usage{Window}Percentile` predicate. */ + private static final java.util.Set WINDOWS = + java.util.Set.of("daily", "weekly", "monthly"); + + private ImportanceQueryBuilder() {} + + public static String build(String entityType, String window, int limit) { + String safeType = validateEntityType(entityType); + String safeWindow = validateWindow(window); + int safeLimit = clamp(limit, MIN_LIMIT, MAX_LIMIT); + String classLocalName = capitalize(safeType); + String pctPredicate = "usage" + capitalize(safeWindow) + "Percentile"; + + return String.join( + "\n", + "PREFIX om: ", + "PREFIX rdfs: ", + "PREFIX prov: ", + "PREFIX xsd: ", + "SELECT ?entity ?fqn ?label ?usagePct ?downstreamCount ?score WHERE {", + " {", + " SELECT ?entity (COUNT(?downstream) AS ?downstreamCount) WHERE {", + " ?entity a om:" + classLocalName + " .", + " OPTIONAL { ?downstream prov:wasDerivedFrom ?entity }", + " } GROUP BY ?entity", + " }", + " ?entity om:fullyQualifiedName ?fqn .", + " OPTIONAL { ?entity rdfs:label ?label }", + " OPTIONAL { ?entity om:" + pctPredicate + " ?usagePct }", + " OPTIONAL { ?entity om:centralityScore ?centrality }", + " {", + " SELECT (MAX(?dc) AS ?maxDownstream) WHERE {", + " SELECT (COUNT(?ds) AS ?dc) WHERE {", + " ?e a om:" + classLocalName + " .", + " OPTIONAL { ?ds prov:wasDerivedFrom ?e }", + " } GROUP BY ?e", + " }", + " }", + " BIND(COALESCE(?usagePct, 0.0) / 100.0 AS ?usageNorm)", + " BIND(", + " IF(?maxDownstream > 0,", + " xsd:double(?downstreamCount) / xsd:double(?maxDownstream),", + " 0.0)", + " AS ?downstreamNorm)", + " BIND(COALESCE(?centrality, 0.0) AS ?centralityNorm)", + " BIND(", + " (0.6 * ?usageNorm) + (0.4 * ?downstreamNorm) + (0.0 * ?centralityNorm)", + " AS ?score)", + "}", + "ORDER BY DESC(?score) DESC(?downstreamCount)", + "LIMIT " + safeLimit); + } + + static String validateEntityType(String entityType) { + if (entityType == null || entityType.isBlank()) { + throw new IllegalArgumentException("'entityType' is required"); + } + String trimmed = entityType.trim(); + if (!trimmed.matches("[a-zA-Z][a-zA-Z0-9]*")) { + throw new IllegalArgumentException( + "'entityType' must be alphanumeric (got '" + entityType + "')"); + } + return trimmed; + } + + static String validateWindow(String window) { + if (window == null || window.isBlank()) { + return "daily"; + } + String lower = window.trim().toLowerCase(java.util.Locale.ROOT); + if (!WINDOWS.contains(lower)) { + throw new IllegalArgumentException( + "'window' must be one of " + WINDOWS + " (got '" + window + "')"); + } + return lower; + } + + static int clamp(int v, int lo, int hi) { + return Math.min(Math.max(v, lo), hi); + } + + private static String capitalize(String s) { + if (s == null || s.isEmpty()) return s; + return Character.toUpperCase(s.charAt(0)) + s.substring(1); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/LineagePathBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/LineagePathBuilder.java new file mode 100644 index 000000000000..ead72dfbd256 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/LineagePathBuilder.java @@ -0,0 +1,167 @@ +package org.openmetadata.service.rdf.insights; + +import java.net.URI; +import java.net.URISyntaxException; +import java.util.Collection; +import java.util.LinkedHashSet; +import java.util.Locale; +import java.util.Set; + +/** + * Builds the SPARQL fragments used by {@link LineagePathFinder} to walk one BFS frontier step + * across the lineage graph. + * + *

The lineage graph is built from three predicates emitted by {@code RdfPropertyMapper}: + * + *

    + *
  • {@code prov:wasDerivedFrom} — entity → its upstream source + *
  • {@code om:upstream} — same direction, OpenMetadata-flavored alias + *
  • {@code om:downstream} — entity → an entity that derives from it + *
+ * + *

For an "upstream" walk from {@code A}, the algorithm follows {@code A prov:wasDerivedFrom ?x} + * and {@code A om:upstream ?x}. For a "downstream" walk it inverts {@code prov:wasDerivedFrom} + * (asks for {@code ?x prov:wasDerivedFrom A}) and follows {@code A om:downstream ?x}. {@code both} + * does both. + * + *

All inputs are validated; any URI not recognized by {@link URI} is rejected before SPARQL is + * emitted. The class is intentionally side-effect-free so it can be exhaustively unit-tested. + */ +public final class LineagePathBuilder { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String PROV_NS = "http://www.w3.org/ns/prov#"; + + /** Default upper bound on BFS depth. */ + public static final int DEFAULT_MAX_HOPS = 6; + + /** Hard cap so a buggy caller can't ask for thousands of frontier expansions. */ + public static final int HARD_MAX_HOPS = 25; + + private LineagePathBuilder() {} + + /** Allowed walk directions for {@link #frontierQuery(Collection, Direction)}. */ + public enum Direction { + UPSTREAM, + DOWNSTREAM, + BOTH; + + public static Direction parse(String value) { + if (value == null || value.isBlank()) return UPSTREAM; + return switch (value.trim().toLowerCase(Locale.ROOT)) { + case "upstream" -> UPSTREAM; + case "downstream" -> DOWNSTREAM; + case "both" -> BOTH; + default -> throw new IllegalArgumentException( + "direction must be one of: upstream, downstream, both (got: " + value + ")"); + }; + } + } + + /** Validate a node URI. Returns the URI string. Throws if missing/blank/malformed. */ + public static String validateNodeUri(String label, String uri) { + if (uri == null || uri.isBlank()) { + throw new IllegalArgumentException(label + " is required"); + } + String trimmed = uri.trim(); + if (trimmed.contains(">") || trimmed.contains("<") || trimmed.contains("\n")) { + throw new IllegalArgumentException(label + " contains illegal characters"); + } + URI parsed; + try { + parsed = new URI(trimmed); + } catch (URISyntaxException e) { + throw new IllegalArgumentException(label + " is not a valid URI: " + e.getMessage()); + } + if (!parsed.isAbsolute()) { + throw new IllegalArgumentException(label + " must be an absolute URI"); + } + String scheme = parsed.getScheme(); + if (!"http".equalsIgnoreCase(scheme) && !"https".equalsIgnoreCase(scheme)) { + throw new IllegalArgumentException( + label + " must use http or https scheme (got: " + scheme + ")"); + } + return trimmed; + } + + /** Clamp maxHops to [1, HARD_MAX_HOPS]; null or < 1 falls back to DEFAULT_MAX_HOPS. */ + public static int clampMaxHops(Integer requested) { + if (requested == null || requested < 1) return DEFAULT_MAX_HOPS; + return Math.min(requested, HARD_MAX_HOPS); + } + + /** + * Build a SPARQL SELECT that, given a frontier of node URIs, returns every + * (?from, ?to, ?predicate) triple reachable in one hop in the requested direction. The result is + * always written so that ?from is the frontier node being expanded — i.e. for a downstream walk + * the inverse of {@code prov:wasDerivedFrom} is rewritten so the caller can treat ?from as + * "current" and ?to as "next" without branching. + */ + public static String frontierQuery(Collection frontier, Direction direction) { + if (frontier == null || frontier.isEmpty()) { + throw new IllegalArgumentException("frontier must contain at least one URI"); + } + Set validated = new LinkedHashSet<>(); + for (String uri : frontier) validated.add(validateNodeUri("frontier node", uri)); + + StringBuilder values = new StringBuilder(); + for (String uri : validated) values.append(" <").append(uri).append(">\n"); + + StringBuilder unions = new StringBuilder(); + if (direction == Direction.UPSTREAM || direction == Direction.BOTH) { + unions.append(union("?from prov:wasDerivedFrom ?to", "prov:wasDerivedFrom")); + unions.append(" UNION\n"); + unions.append(union("?from om:upstream ?to", "om:upstream")); + } + if (direction == Direction.DOWNSTREAM || direction == Direction.BOTH) { + if (unions.length() > 0) unions.append(" UNION\n"); + unions.append(union("?to prov:wasDerivedFrom ?from", "^prov:wasDerivedFrom")); + unions.append(" UNION\n"); + unions.append(union("?from om:downstream ?to", "om:downstream")); + } + + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "PREFIX prov: <" + PROV_NS + ">", + "SELECT ?from ?to ?predicate WHERE {", + " VALUES ?from {", + values.toString().stripTrailing(), + " }", + " {", + unions.toString(), + " }", + " FILTER(?to != ?from)", + "}"); + } + + /** + * Build a SPARQL SELECT that returns the rdf:type values for a set of nodes. Used to decorate + * the final path with class info without round-tripping per node. + */ + public static String typesQuery(Collection nodes) { + if (nodes == null || nodes.isEmpty()) { + throw new IllegalArgumentException("nodes must contain at least one URI"); + } + Set validated = new LinkedHashSet<>(); + for (String uri : nodes) validated.add(validateNodeUri("node", uri)); + + StringBuilder values = new StringBuilder(); + for (String uri : validated) values.append(" <").append(uri).append(">\n"); + + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "SELECT ?node ?type WHERE {", + " VALUES ?node {", + values.toString().stripTrailing(), + " }", + " ?node a ?type .", + " FILTER(STRSTARTS(STR(?type), \"" + OM_NS + "\"))", + "}"); + } + + private static String union(String triple, String predicateLabel) { + return " { " + triple + " . BIND(\"" + predicateLabel + "\" AS ?predicate) }"; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/LineagePathFinder.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/LineagePathFinder.java new file mode 100644 index 000000000000..ea54f876b3bd --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/LineagePathFinder.java @@ -0,0 +1,251 @@ +package org.openmetadata.service.rdf.insights; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Deque; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import lombok.extern.slf4j.Slf4j; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.service.rdf.RdfRepository; + +/** + * BFS path finder over the lineage graph. Walks one frontier at a time using SPARQL queries built + * by {@link LineagePathBuilder}, so the algorithm is independent of dataset size — only the active + * frontier is held in memory. + * + *

Termination conditions: + * + *

    + *
  • Target reached → reconstruct path via parent map; return {@code found = true}. + *
  • Frontier becomes empty → return {@code found = false}. + *
  • maxHops reached → return {@code found = false}. + *
+ * + *

The {@code visited} set guards against cycles. Each node is expanded at most once. The found + * path is always the shortest one (BFS), with ties broken by SPARQL result order. + */ +@Slf4j +public final class LineagePathFinder { + + private final RdfRepository repository; + + public LineagePathFinder(RdfRepository repository) { + this.repository = repository; + } + + /** + * Walk the lineage graph from {@code fromUri} to {@code toUri}. + * + * @param fromUri starting node URI (validated) + * @param toUri target node URI (validated) + * @param direction upstream / downstream / both + * @param maxHops upper bound on path length; clamped via {@link LineagePathBuilder#clampMaxHops} + */ + public Path findPath( + String fromUri, String toUri, LineagePathBuilder.Direction direction, Integer maxHops) { + String from = LineagePathBuilder.validateNodeUri("from", fromUri); + String to = LineagePathBuilder.validateNodeUri("to", toUri); + LineagePathBuilder.Direction dir = + direction == null ? LineagePathBuilder.Direction.UPSTREAM : direction; + int hopBudget = LineagePathBuilder.clampMaxHops(maxHops); + + if (from.equals(to)) { + return Path.found(from, to, dir, hopBudget, List.of(new Hop(0, from, null, List.of()))); + } + + Map parents = new HashMap<>(); + Set visited = new LinkedHashSet<>(); + visited.add(from); + Set frontier = new LinkedHashSet<>(); + frontier.add(from); + + int depth = 0; + String reached = null; + while (!frontier.isEmpty() && depth < hopBudget) { + Map nextLevel = expandFrontier(frontier, dir, visited); + if (nextLevel.isEmpty()) break; + depth++; + + for (Map.Entry e : nextLevel.entrySet()) { + String node = e.getKey(); + if (parents.putIfAbsent(node, e.getValue()) == null) { + visited.add(node); + } + if (node.equals(to)) { + reached = node; + break; + } + } + if (reached != null) break; + frontier = new LinkedHashSet<>(nextLevel.keySet()); + } + + if (reached == null) { + return Path.notFound(from, to, dir, hopBudget); + } + List hops = reconstructPath(from, to, parents); + decorateWithTypes(hops); + return Path.found(from, to, dir, hopBudget, hops); + } + + /** + * Run one frontier query and return a map of (newly discovered node → ParentEdge). Already + * visited nodes are filtered out so the BFS stays acyclic. + */ + Map expandFrontier( + Set frontier, LineagePathBuilder.Direction direction, Set visited) { + String sparql = LineagePathBuilder.frontierQuery(frontier, direction); + String json; + try { + json = repository.executeSparqlQuery(sparql, "application/sparql-results+json"); + } catch (Exception e) { + LOG.warn("Path finder frontier query failed: {}", e.getMessage()); + return Map.of(); + } + return parseFrontierResult(json, visited); + } + + static Map parseFrontierResult(String json, Set visited) { + Map next = new LinkedHashMap<>(); + if (json == null || json.isBlank()) return next; + try { + JsonNode root = JsonUtils.readTree(json); + JsonNode bindings = root.path("results").path("bindings"); + if (!bindings.isArray()) return next; + for (JsonNode row : bindings) { + String from = textValue(row, "from"); + String to = textValue(row, "to"); + String predicate = textValue(row, "predicate"); + if (from == null || to == null || predicate == null) continue; + if (visited.contains(to)) continue; + next.putIfAbsent(to, new ParentEdge(from, predicate)); + } + } catch (Exception e) { + LOG.warn("Failed to parse path frontier result: {}", e.getMessage()); + } + return next; + } + + private List reconstructPath(String from, String to, Map parents) { + Deque stack = new ArrayDeque<>(); + String cursor = to; + int safety = parents.size() + 2; + while (cursor != null && !cursor.equals(from) && safety-- > 0) { + ParentEdge edge = parents.get(cursor); + if (edge == null) break; + stack.push(new Hop(0, cursor, edge.predicate(), List.of())); + cursor = edge.parent(); + } + stack.push(new Hop(0, from, null, List.of())); + + List ordered = new ArrayList<>(stack.size()); + int step = 0; + while (!stack.isEmpty()) { + Hop h = stack.pop(); + ordered.add(new Hop(step++, h.node(), h.predicate(), h.rdfTypes())); + } + return ordered; + } + + /** Single SPARQL round-trip to fetch {@code rdf:type} for every node in the path. */ + void decorateWithTypes(List hops) { + if (hops.isEmpty()) return; + Set nodes = new LinkedHashSet<>(); + for (Hop h : hops) nodes.add(h.node()); + String sparql = LineagePathBuilder.typesQuery(nodes); + String json; + try { + json = repository.executeSparqlQuery(sparql, "application/sparql-results+json"); + } catch (Exception e) { + LOG.debug("Type decoration failed (non-fatal): {}", e.getMessage()); + return; + } + Map> types = parseTypesResult(json); + for (int i = 0; i < hops.size(); i++) { + Hop h = hops.get(i); + List t = types.getOrDefault(h.node(), List.of()); + hops.set(i, new Hop(h.step(), h.node(), h.predicate(), t)); + } + } + + static Map> parseTypesResult(String json) { + Map> result = new HashMap<>(); + if (json == null || json.isBlank()) return result; + try { + JsonNode root = JsonUtils.readTree(json); + JsonNode bindings = root.path("results").path("bindings"); + if (!bindings.isArray()) return result; + for (JsonNode row : bindings) { + String node = textValue(row, "node"); + String type = textValue(row, "type"); + if (node == null || type == null) continue; + result.computeIfAbsent(node, k -> new ArrayList<>()).add(type); + } + for (Map.Entry> e : result.entrySet()) { + Set uniq = new LinkedHashSet<>(e.getValue()); + e.setValue(List.copyOf(uniq)); + } + } catch (Exception e) { + LOG.debug("Failed to parse type decoration result: {}", e.getMessage()); + } + return Collections.unmodifiableMap(result); + } + + private static String textValue(JsonNode row, String varName) { + JsonNode node = row.path(varName); + if (node.isMissingNode() || node.isNull()) return null; + JsonNode value = node.path("value"); + return value.isMissingNode() || value.isNull() ? null : value.asText(); + } + + /** Edge that points back to a node's BFS parent. */ + record ParentEdge(String parent, String predicate) {} + + /** One hop in a returned path. */ + public record Hop(int step, String node, String predicate, List rdfTypes) {} + + /** + * BFS path response. {@code hops} is {@code nodes.size() - 1} when found, else {@code 0}. + */ + public record Path( + String from, + String to, + String direction, + int maxHops, + boolean found, + int hops, + List nodes) { + + static Path notFound(String from, String to, LineagePathBuilder.Direction dir, int maxHops) { + return new Path(from, to, dir.name().toLowerCase(), maxHops, false, 0, List.of()); + } + + static Path found( + String from, String to, LineagePathBuilder.Direction dir, int maxHops, List nodes) { + return new Path( + from, to, dir.name().toLowerCase(), maxHops, true, Math.max(0, nodes.size() - 1), nodes); + } + } + + /** Helper to build a Hop list for unit tests. */ + static List hopList(String... nodes) { + List out = new ArrayList<>(); + for (int i = 0; i < nodes.length; i++) { + out.add(new Hop(i, nodes[i], i == 0 ? null : "prov:wasDerivedFrom", List.of())); + } + return out; + } + + /** No-op visited-set placeholder used in tests. */ + static Set emptyVisited() { + return new HashSet<>(); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/Louvain.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/Louvain.java new file mode 100644 index 000000000000..1d46acb96dc4 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/Louvain.java @@ -0,0 +1,228 @@ +package org.openmetadata.service.rdf.insights; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Modularity-optimizing community detection — the greedy "first pass" of the Louvain algorithm + * (Blondel et al., 2008). Each node starts in its own community; we repeatedly move every node to + * the neighbouring community that yields the highest modularity gain until no move improves the + * partition. The aggregation/recursion step of full multi-level Louvain is intentionally left out — + * a single greedy pass already produces high-quality communities on the kinds of graphs an + * OpenMetadata catalog produces (lineage chains, tag co-occurrence) and keeps the algorithm + * deterministic and ~150 lines instead of ~400. + * + *

Determinism: nodes are processed in input-map iteration order; candidate communities are + * sorted by integer id; ties on modularity gain favour the lower-numbered community. A given input + * graph therefore always produces the same partition, which is exactly what the Phase 3.2 success + * criterion calls for ("Louvain run produces om:Community resources with deterministic + * membership for the seed graph"). + * + *

Edge weight handling: the input is treated as undirected; if both directions are present they + * sum, single-direction edges are mirrored. Self-loops are ignored. Negative weights are clamped + * to zero — modularity isn't well-defined for them. + * + *

Modularity gain when inserting node {@code i} into community {@code C} (i has been removed + * from its current community before the test) is the standard Louvain shortcut: + * + * 

+ *   ΔQ ∝ k_iC − Σ_tot(C) · k_i / (2m)
+ *
+ * + * where {@code k_iC} = sum of weights from i to members of C, {@code Σ_tot(C)} = sum of degrees of + * nodes in C, {@code k_i} = degree of i, {@code 2m} = sum of all degrees. The constant 1/(2m) is + * the same across candidates so we drop it for the argmax. + */ +public final class Louvain { + + private static final int DEFAULT_MAX_ITERATIONS = 32; + + private final int maxIterations; + + public Louvain() { + this(DEFAULT_MAX_ITERATIONS); + } + + public Louvain(int maxIterations) { + if (maxIterations < 1) { + throw new IllegalArgumentException("maxIterations must be >= 1"); + } + this.maxIterations = maxIterations; + } + + /** + * Run greedy modularity optimization on a weighted graph. + * + * @param graph adjacency map. {@code graph.get(a).get(b)} = edge weight from a to b. Undirected; + * symmetrized internally. Null or empty graphs return an empty result. + * @param node id type (must support equals/hashCode and have a stable toString) + * @return final partition + modularity + iteration count + */ + public Result compute(Map> graph) { + if (graph == null || graph.isEmpty()) return new Result<>(Map.of(), 0.0, 0); + + List nodes = new ArrayList<>(graph.keySet()); + int n = nodes.size(); + Map idx = new LinkedHashMap<>(n); + for (int i = 0; i < n; i++) idx.put(nodes.get(i), i); + + Map> adj = new LinkedHashMap<>(n); + for (int i = 0; i < n; i++) adj.put(i, new LinkedHashMap<>()); + addAllEdges(graph, idx, adj); + + double[] degree = new double[n]; + double totalWeight = 0.0; + for (int i = 0; i < n; i++) { + double d = 0.0; + for (double w : adj.get(i).values()) d += w; + degree[i] = d; + totalWeight += d; + } + if (totalWeight == 0.0) { + return singletonsResult(nodes); + } + + int[] community = new int[n]; + double[] commTotal = new double[n]; + for (int i = 0; i < n; i++) { + community[i] = i; + commTotal[i] = degree[i]; + } + + int iterations = 0; + boolean moved = true; + while (moved && iterations < maxIterations) { + moved = false; + iterations++; + for (int i = 0; i < n; i++) { + int chosen = considerMoves(i, community, commTotal, degree, totalWeight, adj); + if (chosen != community[i]) { + commTotal[community[i]] -= degree[i]; + commTotal[chosen] += degree[i]; + community[i] = chosen; + moved = true; + } + } + } + + Map renumbered = renumber(community); + Map finalPartition = new LinkedHashMap<>(n); + for (int i = 0; i < n; i++) finalPartition.put(nodes.get(i), renumbered.get(community[i])); + double modularity = computeModularity(adj, community, degree, totalWeight); + return new Result<>(finalPartition, modularity, iterations); + } + + private static void addAllEdges( + Map> graph, Map idx, Map> adj) { + for (Map.Entry> e : graph.entrySet()) { + Integer src = idx.get(e.getKey()); + if (src == null || e.getValue() == null) continue; + for (Map.Entry e2 : e.getValue().entrySet()) { + Integer dst = idx.get(e2.getKey()); + if (dst == null || dst.equals(src)) continue; + double w = e2.getValue() == null ? 0.0 : Math.max(0.0, e2.getValue()); + if (w == 0.0) continue; + adj.get(src).merge(dst, w, Double::sum); + adj.get(dst).merge(src, w, Double::sum); + } + } + } + + /** Choose the best community for node {@code i}, breaking ties toward lower community id. */ + private static int considerMoves( + int i, + int[] community, + double[] commTotal, + double[] degree, + double totalWeight, + Map> adj) { + int currentComm = community[i]; + double k_i = degree[i]; + + Map commLinks = new LinkedHashMap<>(); + for (Map.Entry e : adj.get(i).entrySet()) { + if (e.getKey() == i) continue; + commLinks.merge(community[e.getKey()], e.getValue(), Double::sum); + } + + commTotal[currentComm] -= k_i; + + int bestComm = currentComm; + double bestGain = 0.0; + List candidates = new ArrayList<>(commLinks.keySet()); + Collections.sort(candidates); + for (int c : candidates) { + double k_iC = commLinks.getOrDefault(c, 0.0); + double gain = k_iC - commTotal[c] * k_i / totalWeight; + if (gain > bestGain || (gain == bestGain && c < bestComm)) { + bestGain = gain; + bestComm = c; + } + } + commTotal[currentComm] += k_i; + return bestComm; + } + + /** Compress community ids to a dense [0..k-1] range, preserving discovery order. */ + static Map renumber(int[] community) { + Map map = new LinkedHashMap<>(); + for (int c : community) { + map.computeIfAbsent(c, k -> map.size()); + } + return map; + } + + /** Modularity Q = (1/2m) Σ [A_ij − k_i k_j / 2m] · δ(c_i, c_j). */ + static double computeModularity( + Map> adj, + int[] community, + double[] degree, + double totalWeight) { + if (totalWeight == 0) return 0.0; + double q = 0.0; + int n = adj.size(); + for (int i = 0; i < n; i++) { + for (Map.Entry e : adj.get(i).entrySet()) { + int j = e.getKey(); + if (community[i] != community[j]) continue; + double aij = e.getValue(); + double expected = degree[i] * degree[j] / totalWeight; + q += aij - expected; + } + } + return q / totalWeight; + } + + private static Result singletonsResult(List nodes) { + Map partition = new LinkedHashMap<>(nodes.size()); + for (int i = 0; i < nodes.size(); i++) partition.put(nodes.get(i), i); + return new Result<>(partition, 0.0, 0); + } + + /** + * Result of a Louvain run. + * + * @param communityByNode every input node mapped to a dense 0-based community id + * @param modularity Q score of the final partition; higher is better, [-1, 1] + * @param iterations number of greedy passes performed + */ + public record Result(Map communityByNode, double modularity, int iterations) { + + /** Inverse view: community id → list of members in input iteration order. */ + public Map> membersByCommunity() { + Map> out = new LinkedHashMap<>(); + for (Map.Entry e : communityByNode.entrySet()) { + out.computeIfAbsent(e.getValue(), k -> new ArrayList<>()).add(e.getKey()); + } + return out; + } + + /** Distinct community count. */ + public int communityCount() { + return (int) communityByNode.values().stream().distinct().count(); + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/PageRank.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/PageRank.java new file mode 100644 index 000000000000..1bbb67ed9ea8 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/PageRank.java @@ -0,0 +1,167 @@ +package org.openmetadata.service.rdf.insights; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +/** + * Iterative weighted PageRank with proper handling of dangling nodes and + * disconnected components. The output is a normalized importance score in + * {@code [0,1]} that sums to 1.0 across all nodes in the graph. + * + *

This is intentionally a tiny, dependency-free implementation. JGraphT or Spark GraphFrames + * would be appropriate for million-edge graphs; OpenMetadata's metadata graph for a single + * tenant typically has < 50k nodes and the simpler approach keeps the codebase lean. + * + *

Algorithm (classic Brin/Page formulation, weighted edges): + * + * 

{@code
+ * for each iteration:
+ *   for each node v:
+ *     newScore[v] = (1 - d) / N
+ *                 + d * sum over u in inEdges(v) of (score[u] * w(u,v) / outWeight(u))
+ *                 + d * danglingMass / N    // redistribute mass from sinks
+ * }
+ * + *

Stops when {@code max |newScore[v] - score[v]| < tolerance} or {@code maxIterations} + * is reached. Output normalizes to sum 1.0. + */ +public final class PageRank { + + /** Standard damping factor. */ + public static final double DEFAULT_DAMPING = 0.85; + + public static final int DEFAULT_MAX_ITERATIONS = 100; + public static final double DEFAULT_TOLERANCE = 1e-6; + + private final double damping; + private final int maxIterations; + private final double tolerance; + + public PageRank() { + this(DEFAULT_DAMPING, DEFAULT_MAX_ITERATIONS, DEFAULT_TOLERANCE); + } + + public PageRank(double damping, int maxIterations, double tolerance) { + if (damping <= 0 || damping >= 1) { + throw new IllegalArgumentException("damping must be in (0, 1)"); + } + if (maxIterations <= 0) { + throw new IllegalArgumentException("maxIterations must be positive"); + } + if (tolerance <= 0) { + throw new IllegalArgumentException("tolerance must be positive"); + } + this.damping = damping; + this.maxIterations = maxIterations; + this.tolerance = tolerance; + } + + /** + * Compute weighted PageRank. + * + * @param outgoing for each node, the map of {@code targetNode → edgeWeight}. Source nodes + * with no outgoing edges (dangling sinks) are still included in the result; their mass is + * redistributed uniformly. Nodes that appear only as targets are treated as having no + * outgoing edges. Empty input returns an empty map. + * @return map of {@code node → score}, scores in [0, 1] summing to 1.0 (modulo floating-point + * drift). Iteration count is reported via {@link Result#iterations()}. + */ + public Result compute(Map> outgoing) { + Objects.requireNonNull(outgoing, "outgoing"); + // Nodes = union of sources and all targets (so a dangling target still gets a score). + Map> graph = new HashMap<>(); + for (Map.Entry> e : outgoing.entrySet()) { + graph.computeIfAbsent(e.getKey(), k -> new HashMap<>()); + for (String t : e.getValue().keySet()) { + graph.computeIfAbsent(t, k -> new HashMap<>()); + } + graph.get(e.getKey()).putAll(e.getValue()); + } + int n = graph.size(); + if (n == 0) { + return new Result<>(Collections.emptyMap(), 0, true); + } + + // Pre-compute outgoing weight totals once per source. + Map outWeight = new HashMap<>(); + for (Map.Entry> e : graph.entrySet()) { + double sum = 0.0; + for (Double w : e.getValue().values()) { + if (w != null && w > 0) sum += w; + } + outWeight.put(e.getKey(), sum); + } + + Map score = new HashMap<>(n); + double init = 1.0 / n; + for (String node : graph.keySet()) score.put(node, init); + + int iterations = 0; + boolean converged = false; + while (iterations < maxIterations) { + iterations++; + // Sum of mass from dangling nodes — to be redistributed uniformly. + double dangling = 0.0; + for (Map.Entry e : score.entrySet()) { + if (outWeight.getOrDefault(e.getKey(), 0.0) <= 0) { + dangling += e.getValue(); + } + } + Map next = new HashMap<>(n); + double danglingShare = damping * dangling / n; + double base = (1 - damping) / n + danglingShare; + for (String node : graph.keySet()) next.put(node, base); + + for (Map.Entry> e : graph.entrySet()) { + String src = e.getKey(); + double srcOut = outWeight.getOrDefault(src, 0.0); + if (srcOut <= 0) continue; + double srcScore = score.get(src); + for (Map.Entry edge : e.getValue().entrySet()) { + double w = edge.getValue() == null ? 0 : edge.getValue(); + if (w <= 0) continue; + double contribution = damping * srcScore * (w / srcOut); + next.merge(edge.getKey(), contribution, Double::sum); + } + } + + double maxDelta = 0.0; + for (Map.Entry e : next.entrySet()) { + double d = Math.abs(e.getValue() - score.getOrDefault(e.getKey(), 0.0)); + if (d > maxDelta) maxDelta = d; + } + score = next; + if (maxDelta < tolerance) { + converged = true; + break; + } + } + + // Normalize so scores sum to 1.0 (guards against floating-point drift). + double total = 0.0; + for (double v : score.values()) total += v; + if (total > 0) { + Map normalized = new HashMap<>(n); + for (Map.Entry e : score.entrySet()) { + normalized.put(e.getKey(), e.getValue() / total); + } + score = normalized; + } + return new Result<>(score, iterations, converged); + } + + /** Visible-for-tests convenience: which nodes appear in the result. */ + public static Set nodes(Map> graph) { + Set all = new java.util.HashSet<>(); + for (Map.Entry> e : graph.entrySet()) { + all.add(e.getKey()); + all.addAll(e.getValue().keySet()); + } + return all; + } + + public record Result(Map scores, int iterations, boolean converged) {} +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/RecommendationsQueryBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/RecommendationsQueryBuilder.java new file mode 100644 index 000000000000..391c0a4667a8 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/insights/RecommendationsQueryBuilder.java @@ -0,0 +1,94 @@ +package org.openmetadata.service.rdf.insights; + +/** + * SPARQL builder for Phase 3.4 dataset recommendations. Given an entity URI, scores every other + * entity by graph-topology similarity along three dimensions: + * + *

    + *
  • Tag overlap — number of distinct {@code om:hasTag} values shared with the seed. + *
  • Glossary overlap — number of distinct {@code om:hasGlossaryTerm} values shared. + *
  • Lineage proximity — number of distinct lineage neighbours both entities share, where + * a "neighbour" is reachable in one hop along {@code om:upstream}, {@code om:downstream}, + * {@code prov:wasDerivedFrom}, or its inverse. + *
+ * + *

The composite score is {@code 1.0·tags + 1.5·glossary + 2.0·lineage}. Lineage proximity is + * weighted highest because it's a stronger structural signal than tag co-occurrence — two tables + * that derive from the same source are tightly coupled, whereas two tables sharing a "PII" tag may + * be entirely unrelated. + * + *

The query is a single SPARQL aggregate with three sub-SELECTs unioned together so the engine + * can compute each dimension once and the outer GROUP BY sums them; this keeps the query plan + * cache-friendly and keeps the result set naturally sparse — entities with zero overlap on any + * dimension never appear in the union and so never reach the score formula. + */ +public final class RecommendationsQueryBuilder { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String PROV_NS = "http://www.w3.org/ns/prov#"; + + /** Default top-N when caller doesn't specify. */ + public static final int DEFAULT_LIMIT = 10; + + /** Hard cap so a buggy caller can't ask for hundreds of recommendations. */ + public static final int MAX_LIMIT = 50; + + static final double WEIGHT_TAG = 1.0; + static final double WEIGHT_GLOSSARY = 1.5; + static final double WEIGHT_LINEAGE = 2.0; + + private RecommendationsQueryBuilder() {} + + /** + * Build the recommendations SPARQL. + * + * @param entityUri seed entity URI; validated as absolute http(s) + * @param limit number of recommendations; clamped to [1, {@link #MAX_LIMIT}] + */ + public static String build(String entityUri, int limit) { + String safeUri = LineagePathBuilder.validateNodeUri("entityUri", entityUri); + int safeLimit = clamp(limit, 1, MAX_LIMIT); + return String.join( + "\n", + "PREFIX om: <" + OM_NS + ">", + "PREFIX prov: <" + PROV_NS + ">", + "SELECT ?candidate", + " (SUM(?t) AS ?tagOverlap)", + " (SUM(?g) AS ?glossaryOverlap)", + " (SUM(?l) AS ?lineageOverlap)", + " ((SUM(?t) * " + WEIGHT_TAG + ")", + " + (SUM(?g) * " + WEIGHT_GLOSSARY + ")", + " + (SUM(?l) * " + WEIGHT_LINEAGE + ") AS ?score)", + "WHERE {", + " {", + " SELECT ?candidate (COUNT(DISTINCT ?tag) AS ?t) (0 AS ?g) (0 AS ?l) WHERE {", + " <" + safeUri + "> om:hasTag ?tag .", + " ?candidate om:hasTag ?tag .", + " FILTER(?candidate != <" + safeUri + ">)", + " } GROUP BY ?candidate", + " } UNION {", + " SELECT ?candidate (0 AS ?t) (COUNT(DISTINCT ?term) AS ?g) (0 AS ?l) WHERE {", + " <" + safeUri + "> om:hasGlossaryTerm ?term .", + " ?candidate om:hasGlossaryTerm ?term .", + " FILTER(?candidate != <" + safeUri + ">)", + " } GROUP BY ?candidate", + " } UNION {", + " SELECT ?candidate (0 AS ?t) (0 AS ?g) (COUNT(DISTINCT ?n) AS ?l) WHERE {", + " <" + + safeUri + + "> (om:upstream|om:downstream|prov:wasDerivedFrom|^prov:wasDerivedFrom) ?n .", + " ?candidate (om:upstream|om:downstream|prov:wasDerivedFrom|^prov:wasDerivedFrom) ?n .", + " FILTER(?candidate != <" + safeUri + ">)", + " } GROUP BY ?candidate", + " }", + "}", + "GROUP BY ?candidate", + "ORDER BY DESC(?score) ?candidate", + "LIMIT " + safeLimit); + } + + private static int clamp(int v, int lo, int hi) { + if (v < lo) return lo; + return Math.min(v, hi); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/JsonLdTranslator.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/JsonLdTranslator.java index 4512e9e469b0..81e0a311db45 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/JsonLdTranslator.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/JsonLdTranslator.java @@ -51,7 +51,9 @@ private void loadContexts() { "governance", "quality", "operations", - "lineage" + "lineage", + "ai", + "automation" }; for (String contextName : contexts) { try { @@ -112,6 +114,7 @@ public ObjectNode toJsonLd(EntityInterface entity) { JsonNode entityJson = objectMapper.valueToTree(entity); Map entityMap = objectMapper.convertValue(entityJson, Map.class); addJsonLdPropertiesToReferences(entityMap); + assignColumnIds(entityMap); String entityType = entity.getEntityReference().getType(); String id = baseUri + "entity/" + entityType + "/" + entity.getId(); @@ -173,6 +176,41 @@ private void addJsonLdPropertiesToReferences(Map map) { } } + /** + * Assign FQN-derived URIs to every Column nested in a Table (or another column) so each Column + * is a first-class named resource. The same URI is minted by column-level lineage so SPARQL can + * traverse from a lineage edge to the column it references. + */ + private void assignColumnIds(Map entityMap) { + Object columnsValue = entityMap.get("columns"); + if (columnsValue instanceof java.util.List) { + for (Object column : (java.util.List) columnsValue) { + if (column instanceof Map) { + assignColumnId((Map) column); + } + } + } + } + + private void assignColumnId(Map column) { + Object fqn = column.get("fullyQualifiedName"); + if (fqn instanceof String && !((String) fqn).isEmpty()) { + String columnUri = RdfUtils.columnUri(baseUri, (String) fqn); + if (columnUri != null) { + column.put("@id", columnUri); + column.put("@type", "om:Column"); + } + } + Object children = column.get("children"); + if (children instanceof java.util.List) { + for (Object child : (java.util.List) children) { + if (child instanceof Map) { + assignColumnId((Map) child); + } + } + } + } + private ObjectNode createSimpleJsonLd(EntityInterface entity) { ObjectNode result = objectMapper.createObjectNode(); @@ -320,10 +358,6 @@ private Object selectContext(String entityType) { "testcaseresult", "testcaseresolutionstatus" -> contextCache.get("quality"); case "ingestionpipeline", - "workflow", - "workflowdefinition", - "workflowinstance", - "workflowinstancestate", "eventsubscription", "kpi", "datainsightchart", @@ -332,6 +366,17 @@ private Object selectContext(String entityType) { "appmarketplacedefinition", "document", "page" -> contextCache.get("operations"); + case "llmmodel", + "aiapplication", + "mcpserver", + "mcpexecution", + "agentexecution", + "prompttemplate" -> contextCache.get("ai"); + case "workflow", + "workflowdefinition", + "workflowinstance", + "workflowinstancestate", + "automation" -> contextCache.get("automation"); default -> contextCache.get("base"); }; } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfActivityMapper.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfActivityMapper.java new file mode 100644 index 000000000000..7bee241a6af5 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfActivityMapper.java @@ -0,0 +1,132 @@ +package org.openmetadata.service.rdf.translator; + +import com.fasterxml.jackson.databind.JsonNode; +import org.apache.jena.datatypes.xsd.XSDDatatype; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.Property; +import org.apache.jena.rdf.model.Resource; +import org.apache.jena.vocabulary.RDF; + +/** + * Emits {@code prov:Activity} resources for OpenMetadata pipeline runs. + * + *

OpenMetadata stores the latest run of a pipeline under {@code Pipeline.pipelineStatus}. We + * surface that as a navigable PROV-O activity tied to the pipeline (as the agent associated + * with the activity) and to the input/output datasets via {@code prov:used} / + * {@code prov:generated}, so SPARQL can answer "who ran what at when, against which datasets, + * with what outcome." + */ +public final class RdfActivityMapper { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String PROV_NS = "http://www.w3.org/ns/prov#"; + + private RdfActivityMapper() {} + + static void emitPipelineActivity( + JsonNode pipelineStatus, + String pipelineFqn, + Resource pipelineResource, + String baseUri, + Model model) { + if (pipelineStatus == null || pipelineStatus.isNull() || !pipelineStatus.isObject()) { + return; + } + if (!pipelineStatus.has("timestamp") || pipelineStatus.get("timestamp").isNull()) { + // PROV activity is meaningless without a startedAtTime equivalent. + return; + } + long startMillis = pipelineStatus.get("timestamp").asLong(); + String activityUri = activityUri(pipelineResource, pipelineFqn, startMillis); + Resource activity = model.createResource(activityUri); + activity.addProperty(RDF.type, model.createResource(PROV_NS + "Activity")); + activity.addProperty(RDF.type, model.createResource(OM_NS + "PipelineExecution")); + + String startedAt = java.time.Instant.ofEpochMilli(startMillis).toString(); + activity.addProperty( + model.createProperty(PROV_NS, "startedAtTime"), + model.createTypedLiteral(startedAt, XSDDatatype.XSDdateTime)); + if (pipelineStatus.has("endTime") && pipelineStatus.get("endTime").isNumber()) { + String endedAt = + java.time.Instant.ofEpochMilli(pipelineStatus.get("endTime").asLong()).toString(); + activity.addProperty( + model.createProperty(PROV_NS, "endedAtTime"), + model.createTypedLiteral(endedAt, XSDDatatype.XSDdateTime)); + } + + if (pipelineStatus.has("executionStatus") && !pipelineStatus.get("executionStatus").isNull()) { + activity.addProperty( + model.createProperty(OM_NS, "executionStatus"), + pipelineStatus.get("executionStatus").asText()); + } + if (pipelineStatus.has("executionId") && !pipelineStatus.get("executionId").isNull()) { + activity.addProperty( + model.createProperty(OM_NS, "executionId"), pipelineStatus.get("executionId").asText()); + } + + // PROV-O: prov:wasInformedBy is an Activity → Activity relation. The pipeline run "was + // informed by" its pipeline definition (the template Activity). Previously this used + // prov:wasGeneratedBy, which has domain prov:Entity and range prov:Activity — inverted. + activity.addProperty(model.createProperty(PROV_NS, "wasInformedBy"), pipelineResource); + pipelineResource.addProperty(model.createProperty(OM_NS, "hasExecution"), activity); + + addAgent(activity, pipelineStatus.get("executedBy"), baseUri, model); + addUsedDatasets(activity, pipelineStatus.get("inputs"), "datasetFQN", "used", baseUri, model); + addUsedDatasets( + activity, pipelineStatus.get("outputs"), "datasetFQN", "generated", baseUri, model); + } + + private static void addAgent( + Resource activity, JsonNode executedBy, String baseUri, Model model) { + if (executedBy == null || executedBy.isNull() || !executedBy.has("id")) { + return; + } + String type = + executedBy.has("type") && !executedBy.get("type").isNull() + ? executedBy.get("type").asText() + : "user"; + // Always mint agent IRIs under the deployment's entity namespace, never the ontology + // namespace. JsonLdTranslator/RdfRepository wire the `om:` prefix to the *ontology* URI + // (https://open-metadata.org/ontology/...), so reading that prefix here would place agent + // resources alongside class definitions and mix ontology + instance data. + Resource agent = + model.createResource(baseUri + "entity/" + type + "/" + executedBy.get("id").asText()); + activity.addProperty(model.createProperty(PROV_NS, "wasAssociatedWith"), agent); + } + + private static void addUsedDatasets( + Resource activity, + JsonNode datasets, + String fqnField, + String predicate, + String baseUri, + Model model) { + if (datasets == null || !datasets.isArray()) { + return; + } + Property prop = model.createProperty(PROV_NS, predicate); + for (JsonNode item : datasets) { + if (!item.isObject() || !item.has(fqnField) || item.get(fqnField).isNull()) { + continue; + } + String fqn = item.get(fqnField).asText(); + // Datasets in Pipeline runs are referenced by FQN (no UUID at this layer); mint a stable + // table URI from the FQN. The triplestore may already contain the table at a UUID-based + // URI; both will participate in queries via the om:fullyQualifiedName literal. + String datasetUri = + baseUri + + "entity/datasetByFqn/" + + java.net.URLEncoder.encode(fqn, java.nio.charset.StandardCharsets.UTF_8); + Resource dataset = model.createResource(datasetUri); + dataset.addProperty(model.createProperty(OM_NS, "fullyQualifiedName"), fqn); + activity.addProperty(prop, dataset); + } + } + + private static String activityUri( + Resource pipelineResource, String pipelineFqn, long startMillis) { + String suffix = pipelineFqn != null ? pipelineFqn : pipelineResource.getURI(); + int hash = (suffix + ":" + startMillis).hashCode(); + return pipelineResource.getURI() + "/run/" + Integer.toHexString(hash); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java index cc6887bac101..43f9f0d696c8 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfPropertyMapper.java @@ -50,8 +50,17 @@ public class RdfPropertyMapper { private static final Set STRUCTURED_PROPERTIES = Set.of("lifeCycle", "customProperties", "extension", "certification"); - // Properties that should be omitted from RDF because they are audit/helper data. - private static final Set IGNORED_PROPERTIES = Set.of("changeDescription", "votes"); + // Properties that should be omitted from RDF because they are audit/helper data, or are + // handled by a dedicated emission step elsewhere in this class (e.g. tableConstraints, which + // requires the parent table FQN to mint constrained-column URIs). + private static final Set IGNORED_PROPERTIES = + Set.of( + "changeDescription", + "votes", + "tableConstraints", + "profile", + "pipelineStatus", + "usageSummary"); // Lineage properties that need special handling private static final Set LINEAGE_PROPERTIES = @@ -84,6 +93,37 @@ public void mapEntityToRdf(EntityInterface entity, Resource entityResource, Mode processContextMappings((Map) context, entityJson, entityResource, model); } + // Table-level constraints need the parent table's FQN to resolve column-name references + // into om:Column URIs, so they're emitted here rather than through the field-mapping loop. + if (entityJson.has("tableConstraints") && entityJson.get("tableConstraints").isArray()) { + emitTableConstraints( + entityJson.get("tableConstraints"), + entity.getFullyQualifiedName(), + entityResource, + model); + } + + // Table profile becomes structured DQV measurements rather than an opaque JSON literal. + if (entityJson.has("profile")) { + RdfQualityMapper.emitTableProfile(entityJson.get("profile"), entityResource, model); + } + + // Pipeline runs surface as prov:Activity resources tied back to inputs and outputs. + if (entityJson.has("pipelineStatus") && !entityJson.get("pipelineStatus").isNull()) { + RdfActivityMapper.emitPipelineActivity( + entityJson.get("pipelineStatus"), + entity.getFullyQualifiedName(), + entityResource, + baseUri, + model); + } + + // Usage summary becomes om:usageDaily/Weekly/MonthlyCount + Percentile triples so the + // /v1/rdf/insights/important endpoint can rank entities by real query usage. + if (entityJson.has("usageSummary") && !entityJson.get("usageSummary").isNull()) { + RdfUsageMapper.emitUsageSummary(entityJson.get("usageSummary"), entityResource, model); + } + // Always add standard properties addStandardProperties(entity, entityResource, model); @@ -200,6 +240,11 @@ private void processFieldMapping( return; } + if ("columns".equals(fieldName) && fieldValue.isArray()) { + emitColumns(fieldValue, entityResource, model); + return; + } + if (mapping instanceof String) { // Simple property mapping: "name": "rdfs:label" addSimpleProperty(entityResource, (String) mapping, fieldValue, model); @@ -965,14 +1010,22 @@ private void addLineageDetails( /** * Adds column-level lineage as structured RDF. Enables SPARQL queries like: "Which columns feed - * into column X" or "What transformation is applied to column Y" + * into column X" or "What transformation is applied to column Y". + * + *

Source and destination columns are emitted as URI references via om:fromColumn / + * om:toColumn so that SPARQL property paths can join them with om:Column resources minted on the + * Table side. The original FQN string is preserved as om:fromColumnFqn / om:toColumnFqn for + * back-compatibility with consumers that match by string FQN. */ private void addColumnLineage( JsonNode columnsLineage, Resource lineageDetailsResource, Model model) { Property hasColumnLineage = model.createProperty(OM_NS, "hasColumnLineage"); + Property fromColumn = model.createProperty(OM_NS, "fromColumn"); + Property toColumn = model.createProperty(OM_NS, "toColumn"); + Property fromColumnFqn = model.createProperty(OM_NS, "fromColumnFqn"); + Property toColumnFqn = model.createProperty(OM_NS, "toColumnFqn"); for (JsonNode colLineage : columnsLineage) { - // Create column lineage resource String colLineageUri = lineageDetailsResource.getURI() + "/columnLineage/" + UUID.randomUUID(); Resource colLineageResource = model.createResource(colLineageUri); @@ -980,21 +1033,16 @@ private void addColumnLineage( lineageDetailsResource.addProperty(hasColumnLineage, colLineageResource); colLineageResource.addProperty(RDF.type, model.createResource(OM_NS + "ColumnLineage")); - // Add source columns if (colLineage.has("fromColumns") && colLineage.get("fromColumns").isArray()) { - Property fromColumnProp = model.createProperty(OM_NS, "fromColumn"); for (JsonNode fromCol : colLineage.get("fromColumns")) { - colLineageResource.addProperty(fromColumnProp, fromCol.asText()); + linkColumn(colLineageResource, fromColumn, fromColumnFqn, fromCol, model); } } - // Add destination column if (colLineage.has("toColumn") && !colLineage.get("toColumn").isNull()) { - colLineageResource.addProperty( - model.createProperty(OM_NS, "toColumn"), colLineage.get("toColumn").asText()); + linkColumn(colLineageResource, toColumn, toColumnFqn, colLineage.get("toColumn"), model); } - // Add transformation function if (colLineage.has("function") && !colLineage.get("function").isNull()) { colLineageResource.addProperty( model.createProperty(OM_NS, "transformFunction"), colLineage.get("function").asText()); @@ -1002,6 +1050,27 @@ private void addColumnLineage( } } + private void linkColumn( + Resource colLineageResource, + Property uriProperty, + Property fqnProperty, + JsonNode columnFqnNode, + Model model) { + String fqn = columnFqnNode.asText(); + if (fqn == null || fqn.isEmpty()) { + return; + } + colLineageResource.addProperty(fqnProperty, fqn); + String columnUri = RdfUtils.columnUri(baseUri, fqn); + if (columnUri == null) { + return; + } + Resource columnResource = model.createResource(columnUri); + columnResource.addProperty(RDF.type, model.createResource(OM_NS + "Column")); + columnResource.addProperty(model.createProperty(OM_NS, "fullyQualifiedName"), fqn); + colLineageResource.addProperty(uriProperty, columnResource); + } + /** * Handles full lineage object (entity + nodes + upstreamEdges + downstreamEdges) */ @@ -1083,6 +1152,218 @@ private void addContainerProperty( } } + /** + * Emit each Column in a Table.columns array as a first-class named resource and link the table + * to it via om:hasColumn. URIs are derived from the Column's FQN so that lineage edges + * (om:fromColumn / om:toColumn) resolve to the same resource. + */ + private void emitColumns(JsonNode columns, Resource tableResource, Model model) { + Property hasColumn = model.createProperty(OM_NS, "hasColumn"); + for (JsonNode column : columns) { + if (!column.isObject() || !column.has("fullyQualifiedName")) { + continue; + } + Resource columnResource = buildColumnResource(column, model); + if (columnResource != null) { + tableResource.addProperty(hasColumn, columnResource); + if (column.has("children") && column.get("children").isArray()) { + emitColumnChildren(column.get("children"), columnResource, model); + } + } + } + } + + /** + * Emit table-level constraints (PRIMARY_KEY, UNIQUE, FOREIGN_KEY, ...) from + * {@code Table.tableConstraints[]} as named om:TableConstraint resources, and project + * back onto the constrained columns. For FOREIGN_KEY, also emit + * {@code om:references } triples so SPARQL queries can + * traverse FK edges directly. + */ + private void emitTableConstraints( + JsonNode constraints, String tableFqn, Resource tableResource, Model model) { + if (tableFqn == null || tableFqn.isEmpty()) { + return; + } + Property hasConstraint = model.createProperty(OM_NS, "hasConstraint"); + Property constraintType = model.createProperty(OM_NS, "constraintType"); + Property hasConstrainedColumn = model.createProperty(OM_NS, "hasConstrainedColumn"); + Property hasReferredColumn = model.createProperty(OM_NS, "hasReferredColumn"); + Property references = model.createProperty(OM_NS, "references"); + Property relationshipType = model.createProperty(OM_NS, "relationshipType"); + Property isUnique = model.createProperty(OM_NS, "isUnique"); + Property isPrimaryKey = model.createProperty(OM_NS, "isPrimaryKey"); + Resource tableConstraintClass = model.createResource(OM_NS + "TableConstraint"); + + int index = 0; + for (JsonNode constraint : constraints) { + if (!constraint.isObject() || !constraint.has("constraintType")) { + index++; + continue; + } + String type = constraint.get("constraintType").asText(); + Resource constraintResource = model.createResource(constraintUri(tableResource, type, index)); + constraintResource.addProperty(RDF.type, tableConstraintClass); + constraintResource.addProperty(constraintType, type); + tableResource.addProperty(hasConstraint, constraintResource); + if (constraint.has("relationshipType") && !constraint.get("relationshipType").isNull()) { + constraintResource.addProperty( + relationshipType, constraint.get("relationshipType").asText()); + } + + java.util.List sourceColumns = + resolveColumns(constraint.get("columns"), tableFqn, model); + for (Resource sourceColumn : sourceColumns) { + constraintResource.addProperty(hasConstrainedColumn, sourceColumn); + if ("PRIMARY_KEY".equals(type)) { + sourceColumn.addProperty(isPrimaryKey, model.createTypedLiteral(true)); + sourceColumn.addProperty( + model.createProperty(OM_NS, "isNullable"), model.createTypedLiteral(false)); + sourceColumn.addProperty(isUnique, model.createTypedLiteral(true)); + } else if ("UNIQUE".equals(type)) { + sourceColumn.addProperty(isUnique, model.createTypedLiteral(true)); + } + } + + if ("FOREIGN_KEY".equals(type)) { + java.util.List referredColumns = + resolveReferredColumns(constraint.get("referredColumns"), model); + for (Resource referred : referredColumns) { + constraintResource.addProperty(hasReferredColumn, referred); + } + // Pair source columns with referred columns positionally so SPARQL can traverse + // om:references directly without going through the + // constraint resource. The pairs are in declared array order. + int pairs = Math.min(sourceColumns.size(), referredColumns.size()); + for (int i = 0; i < pairs; i++) { + sourceColumns.get(i).addProperty(references, referredColumns.get(i)); + } + } + index++; + } + } + + private String constraintUri(Resource tableResource, String type, int index) { + return tableResource.getURI() + "/constraint/" + type + "/" + index; + } + + private java.util.List resolveColumns( + JsonNode columnNames, String tableFqn, Model model) { + java.util.List resolved = new java.util.ArrayList<>(); + if (columnNames == null || !columnNames.isArray()) { + return resolved; + } + for (JsonNode name : columnNames) { + if (!name.isTextual()) { + continue; + } + String columnFqn = tableFqn + "." + name.asText(); + Resource columnResource = ensureColumnResource(columnFqn, model); + if (columnResource != null) { + resolved.add(columnResource); + } + } + return resolved; + } + + private java.util.List resolveReferredColumns(JsonNode referred, Model model) { + java.util.List resolved = new java.util.ArrayList<>(); + if (referred == null || !referred.isArray()) { + return resolved; + } + for (JsonNode fqnNode : referred) { + if (!fqnNode.isTextual()) { + continue; + } + Resource columnResource = ensureColumnResource(fqnNode.asText(), model); + if (columnResource != null) { + resolved.add(columnResource); + } + } + return resolved; + } + + private Resource ensureColumnResource(String columnFqn, Model model) { + String columnUri = RdfUtils.columnUri(baseUri, columnFqn); + if (columnUri == null) { + return null; + } + Resource columnResource = model.createResource(columnUri); + columnResource.addProperty(RDF.type, model.createResource(OM_NS + "Column")); + columnResource.addProperty(model.createProperty(OM_NS, "fullyQualifiedName"), columnFqn); + return columnResource; + } + + private void emitColumnChildren(JsonNode children, Resource parentColumn, Model model) { + Property hasChild = model.createProperty(OM_NS, "hasChildColumn"); + for (JsonNode child : children) { + if (!child.isObject() || !child.has("fullyQualifiedName")) { + continue; + } + Resource childResource = buildColumnResource(child, model); + if (childResource != null) { + parentColumn.addProperty(hasChild, childResource); + if (child.has("children") && child.get("children").isArray()) { + emitColumnChildren(child.get("children"), childResource, model); + } + } + } + } + + private Resource buildColumnResource(JsonNode column, Model model) { + String fqn = column.get("fullyQualifiedName").asText(); + String columnUri = RdfUtils.columnUri(baseUri, fqn); + if (columnUri == null) { + return null; + } + Resource columnResource = model.createResource(columnUri); + columnResource.addProperty(RDF.type, model.createResource(OM_NS + "Column")); + columnResource.addProperty(model.createProperty(OM_NS, "fullyQualifiedName"), fqn); + if (column.has("name") && !column.get("name").isNull()) { + columnResource.addProperty(RDFS.label, column.get("name").asText()); + } + if (column.has("dataType") && !column.get("dataType").isNull()) { + columnResource.addProperty( + model.createProperty(OM_NS, "columnDataType"), column.get("dataType").asText()); + } + if (column.has("description") && !column.get("description").isNull()) { + columnResource.addProperty( + model.createProperty(OM_NS, "columnDescription"), column.get("description").asText()); + } + if (column.has("ordinalPosition") && column.get("ordinalPosition").isNumber()) { + columnResource.addProperty( + model.createProperty(OM_NS, "ordinalPosition"), + model.createTypedLiteral(column.get("ordinalPosition").asInt())); + } + if (column.has("constraint") && !column.get("constraint").isNull()) { + applyColumnConstraint(columnResource, column.get("constraint").asText(), model); + } + if (column.has("profile") && !column.get("profile").isNull()) { + RdfQualityMapper.emitColumnProfile(column.get("profile"), columnResource, model); + } + return columnResource; + } + + private void applyColumnConstraint(Resource columnResource, String constraint, Model model) { + Property isPrimaryKey = model.createProperty(OM_NS, "isPrimaryKey"); + Property isNullable = model.createProperty(OM_NS, "isNullable"); + Property isUnique = model.createProperty(OM_NS, "isUnique"); + switch (constraint) { + case "PRIMARY_KEY" -> { + columnResource.addProperty(isPrimaryKey, model.createTypedLiteral(true)); + columnResource.addProperty(isNullable, model.createTypedLiteral(false)); + columnResource.addProperty(isUnique, model.createTypedLiteral(true)); + } + case "UNIQUE" -> columnResource.addProperty(isUnique, model.createTypedLiteral(true)); + case "NOT_NULL" -> columnResource.addProperty(isNullable, model.createTypedLiteral(false)); + case "NULL" -> columnResource.addProperty(isNullable, model.createTypedLiteral(true)); + default -> { + // Unknown / vendor-specific (DIST_KEY, SORT_KEY, etc.) — fall through; surfaced via + // table-level constraints if relevant. + } + } + } + private void addTypedProperty( Resource resource, String propertyId, JsonNode value, String type, Model model) { Property property = createProperty(propertyId, model); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfQualityMapper.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfQualityMapper.java new file mode 100644 index 000000000000..5efb77f13a76 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfQualityMapper.java @@ -0,0 +1,150 @@ +package org.openmetadata.service.rdf.translator; + +import com.fasterxml.jackson.databind.JsonNode; +import java.net.URLEncoder; +import java.nio.charset.StandardCharsets; +import java.util.Map; +import org.apache.jena.datatypes.xsd.XSDDatatype; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.Property; +import org.apache.jena.rdf.model.Resource; +import org.apache.jena.vocabulary.RDF; + +/** + * Emits {@code dqv:QualityMeasurement} triples from OpenMetadata table and column profiles. + * + *

OpenMetadata historically stored these as opaque JSON literals, which makes them + * unqueryable. This mapper turns each numeric profile field into a navigable measurement so + * SPARQL can answer "tables with completeness < 95%", "columns whose null count exceeded N", + * etc. + */ +public final class RdfQualityMapper { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String DQV_NS = "http://www.w3.org/ns/dqv#"; + private static final String PROV_NS = "http://www.w3.org/ns/prov#"; + + // Maps the JSON field name on Table.profile / Column.profile to the metric URI under om:. + // Only numeric metrics are emitted as dqv:value literals; min/max are skipped because their + // type is polymorphic (string / number / dateTime) and a single dqv:value triple can't express + // that without losing the datatype. + private static final Map TABLE_METRICS = + Map.of( + "rowCount", "RowCountMetric", + "columnCount", "ColumnCountMetric", + "sizeInByte", "SizeInBytesMetric"); + + private static final Map COLUMN_METRICS = buildColumnMetricMap(); + + private static Map buildColumnMetricMap() { + return Map.ofEntries( + Map.entry("valuesCount", "ValuesCountMetric"), + Map.entry("validCount", "ValidCountMetric"), + Map.entry("nullCount", "NullCountMetric"), + Map.entry("nullProportion", "NullProportionMetric"), + Map.entry("missingCount", "MissingCountMetric"), + Map.entry("missingPercentage", "MissingPercentageMetric"), + Map.entry("uniqueCount", "UniqueCountMetric"), + Map.entry("uniqueProportion", "UniqueProportionMetric"), + Map.entry("distinctCount", "DistinctCountMetric"), + Map.entry("distinctProportion", "DistinctProportionMetric"), + Map.entry("duplicateCount", "DuplicateCountMetric"), + Map.entry("mean", "MeanMetric"), + Map.entry("sum", "SumMetric"), + Map.entry("stddev", "StddevMetric"), + Map.entry("variance", "VarianceMetric"), + Map.entry("median", "MedianMetric"), + Map.entry("minLength", "MinLengthMetric"), + Map.entry("maxLength", "MaxLengthMetric")); + } + + private RdfQualityMapper() {} + + static void emitTableProfile(JsonNode profile, Resource tableResource, Model model) { + if (profile == null || profile.isNull() || !profile.isObject()) { + return; + } + String timestamp = readTimestamp(profile); + emitMeasurements(profile, TABLE_METRICS, timestamp, tableResource, model); + } + + static void emitColumnProfile(JsonNode profile, Resource columnResource, Model model) { + if (profile == null || profile.isNull() || !profile.isObject()) { + return; + } + String timestamp = readTimestamp(profile); + emitMeasurements(profile, COLUMN_METRICS, timestamp, columnResource, model); + } + + private static void emitMeasurements( + JsonNode profile, + Map metricMap, + String timestamp, + Resource subjectResource, + Model model) { + Property hasMeasurement = model.createProperty(DQV_NS, "hasQualityMeasurement"); + Property isMeasurementOf = model.createProperty(DQV_NS, "isMeasurementOf"); + Property dqvValue = model.createProperty(DQV_NS, "value"); + Property dqvComputedOn = model.createProperty(DQV_NS, "computedOn"); + Property generatedAtTime = model.createProperty(PROV_NS, "generatedAtTime"); + Resource measurementClass = model.createResource(DQV_NS + "QualityMeasurement"); + + for (Map.Entry entry : metricMap.entrySet()) { + String fieldName = entry.getKey(); + String metricLocalName = entry.getValue(); + JsonNode value = profile.get(fieldName); + if (value == null || value.isNull() || !value.isNumber()) { + continue; + } + String measurementUri = measurementUri(subjectResource.getURI(), metricLocalName, timestamp); + Resource measurement = model.createResource(measurementUri); + measurement.addProperty(RDF.type, measurementClass); + measurement.addProperty(isMeasurementOf, model.createResource(OM_NS + metricLocalName)); + measurement.addProperty(dqvComputedOn, subjectResource); + addNumericValue(measurement, dqvValue, value, model); + if (timestamp != null) { + measurement.addProperty( + generatedAtTime, model.createTypedLiteral(timestamp, XSDDatatype.XSDdateTime)); + } + subjectResource.addProperty(hasMeasurement, measurement); + } + } + + private static void addNumericValue( + Resource measurement, Property dqvValue, JsonNode value, Model model) { + if (value.isInt()) { + measurement.addProperty(dqvValue, model.createTypedLiteral(value.asInt())); + } else if (value.isLong()) { + measurement.addProperty(dqvValue, model.createTypedLiteral(value.asLong())); + } else { + measurement.addProperty(dqvValue, model.createTypedLiteral(value.asDouble())); + } + } + + /** + * Mints a deterministic URI for a dqv:QualityMeasurement so that re-emit overwrites prior + * triples instead of creating orphans. Each (subject, metric, timestamp) tuple maps to exactly + * one URI. Missing timestamps fall back to "latest" so back-to-back emits without a profile + * timestamp are idempotent. + */ + static String measurementUri(String subjectUri, String metricLocalName, String timestamp) { + String slot = timestamp == null || timestamp.isEmpty() ? "latest" : timestamp; + String encodedSlot = URLEncoder.encode(slot, StandardCharsets.UTF_8); + return subjectUri + "/measurement/" + metricLocalName + "/" + encodedSlot; + } + + private static String readTimestamp(JsonNode profile) { + JsonNode ts = profile.get("timestamp"); + if (ts == null || ts.isNull()) { + return null; + } + if (ts.isTextual()) { + return ts.asText(); + } + if (ts.isNumber()) { + // OpenMetadata profiles record the timestamp as epoch millis; convert to ISO-8601. + return java.time.Instant.ofEpochMilli(ts.asLong()).toString(); + } + return null; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfUsageMapper.java b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfUsageMapper.java new file mode 100644 index 000000000000..c88c34092218 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/RdfUsageMapper.java @@ -0,0 +1,64 @@ +package org.openmetadata.service.rdf.translator; + +import com.fasterxml.jackson.databind.JsonNode; +import org.apache.jena.datatypes.xsd.XSDDatatype; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.Property; +import org.apache.jena.rdf.model.Resource; + +/** + * Emits RDF triples from {@code Entity.usageSummary} ({@code UsageDetails}). Surfaces query + * usage as a navigable signal so SPARQL — and the {@code /v1/rdf/insights/important} endpoint + * built on top — can rank entities by how often they're actually queried. + * + *

Triples emitted on the entity (only when the corresponding stat is present and numeric): + * + *

    + *
  • {@code om:usageDailyCount}, {@code om:usageDailyPercentile} + *
  • {@code om:usageWeeklyCount}, {@code om:usageWeeklyPercentile} + *
  • {@code om:usageMonthlyCount}, {@code om:usageMonthlyPercentile} + *
  • {@code om:usageDate} ({@code xsd:date}) + *
+ * + *

Percentile values come from OpenMetadata's usage pipeline as 0–100 floats; we keep them in + * that scale (consumers divide by 100 when blending into a 0–1 importance score). + */ +public final class RdfUsageMapper { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + + private RdfUsageMapper() {} + + public static void emitUsageSummary(JsonNode usage, Resource entityResource, Model model) { + if (usage == null || usage.isNull() || !usage.isObject()) { + return; + } + emitStats(usage.get("dailyStats"), "Daily", entityResource, model); + emitStats(usage.get("weeklyStats"), "Weekly", entityResource, model); + emitStats(usage.get("monthlyStats"), "Monthly", entityResource, model); + + JsonNode date = usage.get("date"); + if (date != null && date.isTextual()) { + Property usageDate = model.createProperty(OM_NS, "usageDate"); + entityResource.addProperty( + usageDate, model.createTypedLiteral(date.asText(), XSDDatatype.XSDdate)); + } + } + + private static void emitStats( + JsonNode stats, String window, Resource entityResource, Model model) { + if (stats == null || stats.isNull() || !stats.isObject()) { + return; + } + JsonNode count = stats.get("count"); + if (count != null && count.isNumber()) { + Property countProp = model.createProperty(OM_NS, "usage" + window + "Count"); + entityResource.addProperty(countProp, model.createTypedLiteral(count.asLong())); + } + JsonNode pct = stats.get("percentileRank"); + if (pct != null && pct.isNumber()) { + Property pctProp = model.createProperty(OM_NS, "usage" + window + "Percentile"); + entityResource.addProperty(pctProp, model.createTypedLiteral(pct.asDouble())); + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/mcp/McpUsageResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/mcp/McpUsageResource.java deleted file mode 100644 index 642a9252bfcb..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/mcp/McpUsageResource.java +++ /dev/null @@ -1,366 +0,0 @@ -/* - * Copyright 2025 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.service.resources.mcp; - -import io.swagger.v3.oas.annotations.Operation; -import io.swagger.v3.oas.annotations.Parameter; -import io.swagger.v3.oas.annotations.media.Content; -import io.swagger.v3.oas.annotations.responses.ApiResponse; -import io.swagger.v3.oas.annotations.tags.Tag; -import jakarta.ws.rs.GET; -import jakarta.ws.rs.Path; -import jakarta.ws.rs.Produces; -import jakarta.ws.rs.QueryParam; -import jakarta.ws.rs.core.Context; -import jakarta.ws.rs.core.MediaType; -import jakarta.ws.rs.core.Response; -import jakarta.ws.rs.core.SecurityContext; -import java.time.Duration; -import java.time.Instant; -import java.time.LocalDate; -import java.time.ZoneOffset; -import java.util.LinkedHashMap; -import java.util.LinkedHashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.TreeMap; -import java.util.concurrent.atomic.AtomicLong; -import java.util.function.Consumer; -import java.util.function.Function; -import org.openmetadata.schema.entity.app.App; -import org.openmetadata.schema.entity.app.AppExtension; -import org.openmetadata.schema.entity.app.mcp.McpToolCallUsage; -import org.openmetadata.service.apps.AbstractNativeApplication; -import org.openmetadata.service.apps.ApplicationContext; -import org.openmetadata.service.apps.bundles.mcp.McpAppConstants; -import org.openmetadata.service.jdbi3.AppRepository; -import org.openmetadata.service.resources.Collection; -import org.openmetadata.service.security.Authorizer; - -/** - * Read-only API for MCP tool-call usage. Backed by the {@code apps_extension_time_series} table - * reading from the {@code limits} extension scoped to {@code appName='McpApplication'} — same - * per-app usage bucket CollateAI writes to, isolated by appName. Counts only. No billing, no - * rate-limiting. - */ -@Path("/v1/mcp/usage") -@Tag(name = "MCP Usage", description = "MCP tool-call usage counters and breakdowns.") -@Produces(MediaType.APPLICATION_JSON) -@Collection(name = "mcpUsage") -public class McpUsageResource { - - /** - * Suffixes used to identify bot principals so they can be excluded from unique-user counts and - * per-user breakdowns. Covers both the PascalCase app bot pattern (e.g. {@code - * McpApplicationBot}) and OpenMetadata's lowercase-kebab bot pattern (e.g. {@code - * ingestion-bot}, {@code profiler-bot}, {@code metadata-bot}). Word-boundary aware to avoid - * false positives like {@code robot}. - */ - static final String BOT_SUFFIX_PASCAL = "Bot"; - - static final String BOT_SUFFIX_KEBAB = "-bot"; - static final long DEFAULT_WINDOW_DAYS = 30L; - private static final int PAGE_SIZE = 1000; - - private final Authorizer authorizer; - private final AppRepository appRepository; - - public McpUsageResource(Authorizer authorizer) { - this.authorizer = authorizer; - this.appRepository = new AppRepository(); - } - - @GET - @Path("/summary") - @Operation( - operationId = "getMcpUsageSummary", - summary = "Get aggregate MCP usage counters", - description = - "Returns total/success/failed counts and unique-user count for the supplied window." - + " Defaults to the last 30 days. Admin access required.", - responses = { - @ApiResponse( - responseCode = "200", - description = "Aggregate counters", - content = @Content(mediaType = "application/json")), - @ApiResponse(responseCode = "403", description = "Forbidden. Admin only.") - }) - public Response getSummary( - @Context SecurityContext securityContext, - @Parameter(description = "Window start (epoch millis). Defaults to 30 days ago.") - @QueryParam("startTs") - Long startTs, - @Parameter(description = "Window end exclusive (epoch millis). Defaults to now.") - @QueryParam("endTs") - Long endTs) { - authorizer.authorizeAdmin(securityContext); - long from = resolveStart(startTs); - long to = resolveEnd(endTs); - Response invalid = validateWindow(from, to); - if (invalid != null) { - return invalid; - } - return Response.ok(buildSummary(from, to)).build(); - } - - @GET - @Path("/history") - @Operation( - operationId = "getMcpUsageHistory", - summary = "Daily MCP usage counts", - description = - "Returns a map of UTC start-of-day (epoch millis) to call count. Empty days are filled" - + " with 0 so the series is continuous. Admin access required.") - public Response getHistory( - @Context SecurityContext securityContext, - @QueryParam("startTs") Long startTs, - @QueryParam("endTs") Long endTs) { - authorizer.authorizeAdmin(securityContext); - long from = resolveStart(startTs); - long to = resolveEnd(endTs); - Response invalid = validateWindow(from, to); - if (invalid != null) { - return invalid; - } - return Response.ok(buildDailyHistory(from, to)).build(); - } - - @GET - @Path("/breakdown/tools") - @Operation( - operationId = "getMcpUsageByTool", - summary = "Per-tool call counts", - description = "Counts MCP requests per tool name in the supplied window. Admin only.") - public Response getByTool( - @Context SecurityContext securityContext, - @QueryParam("startTs") Long startTs, - @QueryParam("endTs") Long endTs) { - authorizer.authorizeAdmin(securityContext); - long from = resolveStart(startTs); - long to = resolveEnd(endTs); - Response invalid = validateWindow(from, to); - if (invalid != null) { - return invalid; - } - Map counts = groupByCount(from, to, McpToolCallUsage::getToolName, true); - return Response.ok(counts).build(); - } - - @GET - @Path("/breakdown/users") - @Operation( - operationId = "getMcpUsageByUser", - summary = "Per-user call counts", - description = - "Counts MCP requests per principal in the supplied window. Bot principals" - + " (suffix 'Bot') are excluded. Admin only.") - public Response getByUser( - @Context SecurityContext securityContext, - @QueryParam("startTs") Long startTs, - @QueryParam("endTs") Long endTs) { - authorizer.authorizeAdmin(securityContext); - long from = resolveStart(startTs); - long to = resolveEnd(endTs); - Response invalid = validateWindow(from, to); - if (invalid != null) { - return invalid; - } - Map counts = groupByCount(from, to, McpToolCallUsage::getUserName, false); - return Response.ok(counts).build(); - } - - @GET - @Path("/me") - @Operation( - operationId = "getMcpUsageForMe", - summary = "Self-service MCP usage counters", - description = - "Returns the calling user's total MCP call count and per-tool breakdown for the" - + " supplied window. Any authenticated user.") - public Response getMine( - @Context SecurityContext securityContext, - @QueryParam("startTs") Long startTs, - @QueryParam("endTs") Long endTs) { - String me = securityContext.getUserPrincipal().getName(); - long from = resolveStart(startTs); - long to = resolveEnd(endTs); - Response invalid = validateWindow(from, to); - if (invalid != null) { - return invalid; - } - return Response.ok(buildSelf(me, from, to)).build(); - } - - private Map buildSummary(long from, long to) { - AtomicLong total = new AtomicLong(); - AtomicLong success = new AtomicLong(); - Set users = new LinkedHashSet<>(); - forEachRow( - from, - to, - usage -> { - total.incrementAndGet(); - if (Boolean.TRUE.equals(usage.getSuccess())) { - success.incrementAndGet(); - } - if (usage.getUserName() != null && !isBot(usage.getUserName())) { - users.add(usage.getUserName()); - } - }); - Map body = new LinkedHashMap<>(); - body.put("total", total.get()); - body.put("totalSuccess", success.get()); - body.put("totalFailed", total.get() - success.get()); - body.put("uniqueUsers", users.size()); - body.put("startTs", from); - body.put("endTs", to); - return body; - } - - private Map buildDailyHistory(long from, long to) { - Map daily = new TreeMap<>(); - seedEmptyDays(daily, from, to); - forEachRow(from, to, usage -> daily.merge(startOfDay(usage.getTimestamp()), 1L, Long::sum)); - return daily; - } - - private Map buildSelf(String userName, long from, long to) { - AtomicLong total = new AtomicLong(); - Map byTool = new LinkedHashMap<>(); - forEachRow( - from, - to, - usage -> { - if (!userName.equals(usage.getUserName())) { - return; - } - total.incrementAndGet(); - if (usage.getToolName() != null) { - byTool.merge(usage.getToolName(), 1L, Long::sum); - } - }); - Map body = new LinkedHashMap<>(); - body.put("total", total.get()); - body.put("byTool", byTool); - body.put("startTs", from); - body.put("endTs", to); - return body; - } - - private Map groupByCount( - long from, long to, Function classifier, boolean includeBots) { - Map counts = new LinkedHashMap<>(); - forEachRow( - from, - to, - usage -> { - String key = classifier.apply(usage); - if (key == null) { - return; - } - if (!includeBots && isBot(key)) { - return; - } - counts.merge(key, 1L, Long::sum); - }); - return counts; - } - - /** - * Pages through rows in the half-open window {@code [from, to)} using the upper-bounded SQL - * helper. The SQL filter pins the result set so OFFSET pagination stays consistent across pages - * even if new MCP tool calls are recorded mid-request, preventing duplicate or skipped rows. - */ - private void forEachRow(long from, long to, Consumer visit) { - App app = resolveMcpApp(); - if (app == null) { - return; - } - int offset = 0; - while (true) { - List page = - appRepository.listAppExtensionInWindowByName( - app, - from, - to, - PAGE_SIZE, - offset, - McpToolCallUsage.class, - AppExtension.ExtensionType.LIMITS); - if (page.isEmpty()) { - return; - } - page.forEach(visit); - if (page.size() < PAGE_SIZE) { - return; - } - offset += PAGE_SIZE; - } - } - - private App resolveMcpApp() { - AbstractNativeApplication app = - ApplicationContext.getInstance().getAppIfExists(McpAppConstants.MCP_APP_NAME); - return app != null ? app.getApp() : null; - } - - static long resolveStart(Long startTs) { - if (startTs != null) { - return startTs; - } - return Instant.now().minus(Duration.ofDays(DEFAULT_WINDOW_DAYS)).toEpochMilli(); - } - - static long resolveEnd(Long endTs) { - return endTs != null ? endTs : Instant.now().toEpochMilli(); - } - - /** - * Returns a 400 Response if the resolved window is empty or reversed, otherwise {@code null}. - * Endpoints invoke this before aggregation so callers get an explicit error rather than a - * silently empty payload when they pass a bogus {@code startTs >= endTs}. - */ - static Response validateWindow(long from, long to) { - if (from >= to) { - return Response.status(Response.Status.BAD_REQUEST) - .entity(Map.of("error", "startTs must be before endTs")) - .build(); - } - return null; - } - - static long startOfDay(long epochMillis) { - return LocalDate.ofInstant(Instant.ofEpochMilli(epochMillis), ZoneOffset.UTC) - .atStartOfDay(ZoneOffset.UTC) - .toInstant() - .toEpochMilli(); - } - - static boolean isBot(String principal) { - if (principal == null) { - return false; - } - return principal.endsWith(BOT_SUFFIX_PASCAL) || principal.endsWith(BOT_SUFFIX_KEBAB); - } - - private static void seedEmptyDays(Map daily, long from, long to) { - long cursor = startOfDay(from); - long lastDay = startOfDay(to - 1); - while (cursor <= lastDay) { - daily.put(cursor, 0L); - cursor = cursor + Duration.ofDays(1).toMillis(); - } - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/OntologyDocument.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/OntologyDocument.java new file mode 100644 index 000000000000..628638e3644f --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/OntologyDocument.java @@ -0,0 +1,124 @@ +package org.openmetadata.service.resources.rdf; + +import jakarta.ws.rs.core.MediaType; +import jakarta.ws.rs.core.Response; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import lombok.extern.slf4j.Slf4j; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.ModelFactory; +import org.apache.jena.riot.Lang; +import org.apache.jena.riot.RDFDataMgr; +import org.apache.jena.riot.RDFFormat; + +/** + * Loads the canonical OpenMetadata ontology from the classpath and serves it in the requested + * RDF serialization. The TTL files ship in the openmetadata-spec module: + * + *

    + *
  • {@code /rdf/ontology/openmetadata.ttl} — the main OWL ontology + *
  • {@code /rdf/ontology/openmetadata-prov.ttl} — PROV-aligned extension + *
+ * + * Both are merged into a single response so that consumers don't have to follow imports. + */ +@Slf4j +public final class OntologyDocument { + + private static final String MAIN_RESOURCE = "/rdf/ontology/openmetadata.ttl"; + private static final String PROV_RESOURCE = "/rdf/ontology/openmetadata-prov.ttl"; + + private OntologyDocument() {} + + /** Holder pattern for thread-safe lazy parsing. The merged model is immutable post-load. */ + private static final class Holder { + private static final Model MODEL = loadModel(); + } + + private static Model loadModel() { + Model model = ModelFactory.createDefaultModel(); + readInto(model, MAIN_RESOURCE); + readInto(model, PROV_RESOURCE); + return model; + } + + private static void readInto(Model model, String resourcePath) { + try (InputStream is = OntologyDocument.class.getResourceAsStream(resourcePath)) { + if (is == null) { + LOG.warn("Ontology resource not found on classpath: {}", resourcePath); + return; + } + RDFDataMgr.read(model, is, Lang.TURTLE); + } catch (IOException e) { + LOG.warn("Failed to read ontology resource {}: {}", resourcePath, e.getMessage()); + } catch (RuntimeException e) { + // RDFDataMgr.read can throw Jena RuntimeExceptions (e.g. RiotException) on a malformed + // TTL. Catch broadly so a corrupt ontology file degrades to a partial/empty model rather + // than failing class initialization and taking down /rdf/ontology + MCP describe. + LOG.warn("Failed to parse ontology resource {}: {}", resourcePath, e.getMessage()); + } + } + + /** + * Serialize the merged ontology in the requested RDF format. Returns body, the chosen media + * type, and the file extension. Suitable for callers that don't speak JAX-RS Response (e.g. + * MCP tools). + */ + public static SerializedOntology serializeAsString(String format) { + Format f = Format.parse(format); + ByteArrayOutputStream out = new ByteArrayOutputStream(); + RDFDataMgr.write(out, Holder.MODEL, f.rdfFormat); + return new SerializedOntology(out.toString(StandardCharsets.UTF_8), f.mediaType, f.extension); + } + + public record SerializedOntology(String body, String mediaType, String extension) {} + + static Response serve(String format) { + Format f = Format.parse(format); + try { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + RDFDataMgr.write(out, Holder.MODEL, f.rdfFormat); + return Response.ok(out.toString(StandardCharsets.UTF_8)) + .type(f.mediaType) + .header("Content-Disposition", "inline; filename=openmetadata-ontology." + f.extension) + .build(); + } catch (Exception e) { + LOG.error("Failed to serialize ontology as {}", format, e); + return Response.serverError() + .entity("{\"error\": \"failed to serialize ontology\"}") + .type(MediaType.APPLICATION_JSON) + .build(); + } + } + + private enum Format { + TURTLE(RDFFormat.TURTLE_PRETTY, "text/turtle", "ttl"), + RDFXML(RDFFormat.RDFXML_PRETTY, "application/rdf+xml", "rdf"), + NTRIPLES(RDFFormat.NTRIPLES, "application/n-triples", "nt"), + JSONLD(RDFFormat.JSONLD_PRETTY, "application/ld+json", "jsonld"); + + final RDFFormat rdfFormat; + final String mediaType; + final String extension; + + Format(RDFFormat rdfFormat, String mediaType, String extension) { + this.rdfFormat = rdfFormat; + this.mediaType = mediaType; + this.extension = extension; + } + + static Format parse(String requested) { + if (requested == null) { + return TURTLE; + } + return switch (requested.toLowerCase()) { + case "rdfxml", "rdf+xml", "rdf/xml" -> RDFXML; + case "ntriples", "n-triples" -> NTRIPLES; + case "jsonld", "json-ld", "ld+json" -> JSONLD; + default -> TURTLE; + }; + } + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java index a07e26176622..f73c8039449d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfResource.java @@ -19,19 +19,44 @@ import jakarta.ws.rs.core.Response; import jakarta.ws.rs.core.SecurityContext; import jakarta.ws.rs.core.UriInfo; +import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.StringReader; +import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.LinkedHashSet; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.UUID; import javax.validation.constraints.NotEmpty; import lombok.extern.slf4j.Slf4j; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.ModelFactory; +import org.apache.jena.riot.Lang; +import org.apache.jena.riot.RDFDataMgr; +import org.apache.jena.riot.RDFFormat; +import org.apache.jena.shacl.ValidationReport; +import org.openmetadata.schema.api.configuration.rdf.CustomOntology; +import org.openmetadata.schema.api.configuration.rdf.InferenceRule; import org.openmetadata.schema.api.rdf.SparqlQuery; import org.openmetadata.schema.utils.JsonUtils; import org.openmetadata.service.Entity; import org.openmetadata.service.OpenMetadataApplicationConfig; +import org.openmetadata.service.rdf.RdfIriValidator; import org.openmetadata.service.rdf.RdfRepository; +import org.openmetadata.service.rdf.extension.CustomOntologyRegistry; +import org.openmetadata.service.rdf.extension.CustomOntologyValidator; +import org.openmetadata.service.rdf.federation.SparqlFederationGuard; +import org.openmetadata.service.rdf.inference.InferenceRuleRegistry; +import org.openmetadata.service.rdf.inference.InferenceRuleValidator; +import org.openmetadata.service.rdf.insights.CentralityComputation; +import org.openmetadata.service.rdf.insights.CoOccurrenceQueryBuilder; +import org.openmetadata.service.rdf.insights.CommunityComputation; +import org.openmetadata.service.rdf.insights.ImportanceQueryBuilder; +import org.openmetadata.service.rdf.insights.LineagePathBuilder; +import org.openmetadata.service.rdf.insights.LineagePathFinder; +import org.openmetadata.service.rdf.insights.RecommendationsQueryBuilder; import org.openmetadata.service.rdf.semantic.SemanticSearchEngine; import org.openmetadata.service.resources.Collection; import org.openmetadata.service.security.Authorizer; @@ -49,6 +74,7 @@ public class RdfResource { private volatile RdfRepository rdfRepository; private final Authorizer authorizer; private volatile SemanticSearchEngine semanticSearchEngine; + private volatile SparqlFederationGuard federationGuard; private OpenMetadataApplicationConfig config; public static final String RDF_XML = "application/rdf+xml"; @@ -91,6 +117,15 @@ public void initialize(OpenMetadataApplicationConfig config) { || !Boolean.TRUE.equals(config.getRdfConfiguration().getEnabled())) { LOG.info("RDF support is disabled in configuration"); } + this.federationGuard = new SparqlFederationGuard(config.getRdfConfiguration()); + } + + private synchronized SparqlFederationGuard getFederationGuard() { + if (federationGuard == null) { + // Tests or restarted resource without initialize(); default to closed allowlist. + federationGuard = new SparqlFederationGuard(null); + } + return federationGuard; } @GET @@ -132,6 +167,598 @@ public Response getRdfStatus(@Context SecurityContext securityContext) { return Response.ok().entity(statusJson).type(MediaType.APPLICATION_JSON).build(); } + @POST + @Path("/validate") + @Produces({TURTLE, JSON_LD, MediaType.APPLICATION_JSON}) + @Operation( + operationId = "validateGraph", + summary = "Run SHACL validation against the OpenMetadata knowledge graph", + description = + "Loads the canonical SHACL shapes (rdf/shapes/openmetadata-shapes.ttl) and validates either a single entity's subgraph or the entire dataset against them. The endpoint reports violations; it does not mutate the graph or block writes.", + responses = { + @ApiResponse( + responseCode = "200", + description = + "SHACL validation report (sh:ValidationReport). Conforms field is true when there are no violations.", + content = {@Content(mediaType = TURTLE), @Content(mediaType = JSON_LD)}), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response validateGraph( + @Context SecurityContext securityContext, + @Parameter( + description = + "Optional. Full URI of the entity to scope the validation to (DESCRIBE ). Omit to validate the whole dataset (admin-only, expensive).") + @QueryParam("entityUri") + String entityUri, + @Parameter(description = "Report serialization: turtle (default) or jsonld") + @QueryParam("format") + @DefaultValue("turtle") + String format) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity("{\"error\": \"RDF service not enabled\"}") + .type(MediaType.APPLICATION_JSON) + .build(); + } + + String constructQuery; + if (entityUri != null && !entityUri.isBlank()) { + String validated = RdfIriValidator.validateEntityIri(entityUri); + if (validated == null) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse("entityUri must be an absolute http(s) IRI")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + constructQuery = String.format("DESCRIBE <%s>", validated); + } else { + constructQuery = "CONSTRUCT { ?s ?p ?o } WHERE { ?s ?p ?o }"; + } + + String dataTurtle = getRdfRepository().executeSparqlQueryDirect(constructQuery, "text/turtle"); + Model dataModel = ModelFactory.createDefaultModel(); + try (StringReader reader = new StringReader(dataTurtle)) { + RDFDataMgr.read(dataModel, reader, getRdfRepository().getBaseUri(), Lang.TURTLE); + } catch (Exception e) { + LOG.error("Failed to parse subgraph for SHACL validation", e); + return Response.serverError() + .entity("{\"error\": \"failed to load subgraph for validation\"}") + .type(MediaType.APPLICATION_JSON) + .build(); + } + + ValidationReport report = RdfShaclValidator.validate(dataModel); + + RDFFormat rdfFormat = + "jsonld".equalsIgnoreCase(format) ? RDFFormat.JSONLD_PRETTY : RDFFormat.TURTLE_PRETTY; + String responseMediaType = "jsonld".equalsIgnoreCase(format) ? JSON_LD : TURTLE; + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + RDFDataMgr.write(out, report.getModel(), rdfFormat); + return Response.ok(out.toString(StandardCharsets.UTF_8)) + .type(responseMediaType) + .header("OM-SHACL-Conforms", String.valueOf(report.conforms())) + .build(); + } + + @GET + @Path("/rules") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "listInferenceRules", + summary = "List loaded inference rules", + description = + "Returns all inference rules loaded into this server, in execution order (priority then name). Includes the shipped starter pack plus any rules that have been upserted at runtime.", + responses = { + @ApiResponse(responseCode = "200", description = "List of inference rules"), + @ApiResponse(responseCode = "403", description = "Forbidden") + }) + public Response listInferenceRules(@Context SecurityContext securityContext) { + authorizer.authorizeAdmin(securityContext); + List rules = InferenceRuleRegistry.getInstance().list(); + return Response.ok(JsonUtils.pojoToJson(Map.of("rules", rules))).build(); + } + + @GET + @Path("/rules/{name}") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "getInferenceRule", + summary = "Get a single inference rule by name", + responses = { + @ApiResponse(responseCode = "200", description = "The rule"), + @ApiResponse(responseCode = "403", description = "Forbidden"), + @ApiResponse(responseCode = "404", description = "Rule not found") + }) + public Response getInferenceRule( + @Context SecurityContext securityContext, @PathParam("name") String name) { + authorizer.authorizeAdmin(securityContext); + return InferenceRuleRegistry.getInstance() + .get(name) + .map(rule -> Response.ok(JsonUtils.pojoToJson(rule)).build()) + .orElse( + Response.status(Response.Status.NOT_FOUND) + .entity(buildErrorResponse("Inference rule not found: " + name)) + .type(MediaType.APPLICATION_JSON) + .build()); + } + + @GET + @Path("/insights/important") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "listImportantEntities", + summary = "Rank entities by an importance score that blends usage data and lineage topology", + description = + "Returns the top-N entities of the given type ranked by a composite importance score. The score blends OpenMetadata's existing usage percentile (real query data — 0.6 weight) with downstream lineage edge count (graph topology — 0.4 weight). Once Phase 3.1.b ships, an om:centralityScore from PageRank will fill in for entities that have no query usage data. Results are SPARQL JSON.", + responses = { + @ApiResponse( + responseCode = "200", + description = + "Ranked list of entities with usage percentile, downstream count, and composite score"), + @ApiResponse(responseCode = "400", description = "Invalid entityType, window, or limit"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response listImportantEntities( + @Context SecurityContext securityContext, + @Parameter( + description = + "Entity type to rank (singular: table, dashboard, pipeline, mlmodel, ...). Required.", + required = true) + @QueryParam("entityType") + @NotEmpty + String entityType, + @Parameter(description = "Usage window: daily, weekly, or monthly. Defaults to daily.") + @QueryParam("window") + @DefaultValue("daily") + String window, + @Parameter(description = "Number of results. 1–100, defaults to 20.") + @QueryParam("limit") + @DefaultValue("20") + int limit) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + String sparql; + try { + sparql = ImportanceQueryBuilder.build(entityType, window, limit); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .type(MediaType.APPLICATION_JSON) + .build(); + } + return executeSparqlQuery(sparql, "json", "none"); + } + + @POST + @Path("/insights/recompute-centrality") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "recomputeCentrality", + summary = "Run weighted PageRank on the entity graph and persist scores", + description = + "Triggers Phase 3.1.b's centrality computation: walks lineage / tagging / containment edges of the requested entity type, runs weighted PageRank, and writes the results to the named graph . The /v1/rdf/insights/important endpoint blends these scores in for entities without query usage data. Admin-only; expensive — designed to run on a schedule, but exposed for manual triggering.", + responses = { + @ApiResponse(responseCode = "200", description = "Centrality computation result"), + @ApiResponse(responseCode = "400", description = "Invalid entityType"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response recomputeCentrality( + @Context SecurityContext securityContext, + @Parameter( + description = "Entity type to score (e.g. table, dashboard, pipeline). Required.", + required = true) + @QueryParam("entityType") + @NotEmpty + String entityType) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + try { + CentralityComputation.Result result = + new CentralityComputation(getRdfRepository()).computeAndPersist(entityType); + return Response.ok(JsonUtils.pojoToJson(result)).build(); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .type(MediaType.APPLICATION_JSON) + .build(); + } + } + + @POST + @Path("/insights/recompute-communities") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "recomputeCommunities", + summary = "Run Louvain community detection and persist communities", + description = + "Phase 3.2: extracts the lineage or tag-co-occurrence graph for the requested entity type, runs Louvain modularity optimization, and persists communities to the named graph . Each community is an om:Community resource with om:hasMember triples and modularity score. Admin-only; designed to be triggered on a schedule.", + responses = { + @ApiResponse(responseCode = "200", description = "Community detection result"), + @ApiResponse(responseCode = "400", description = "Invalid entityType or graphType"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response recomputeCommunities( + @Context SecurityContext securityContext, + @Parameter(description = "Entity type to cluster (e.g. table, dashboard).", required = true) + @QueryParam("entityType") + @NotEmpty + String entityType, + @Parameter(description = "Source graph: lineage (default) or tagCoOccurrence.") + @QueryParam("graphType") + @DefaultValue("lineage") + String graphType) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + try { + CommunityComputation.Result result = + new CommunityComputation(getRdfRepository()).computeAndPersist(entityType, graphType); + return Response.ok(JsonUtils.pojoToJson(result)).build(); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .type(MediaType.APPLICATION_JSON) + .build(); + } + } + + @GET + @Path("/insights/communities") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "listCommunities", + summary = "List communities discovered by the latest community-detection run", + description = + "Returns a SPARQL SELECT JSON document with rows of (community, size, modularity, member) for the named graph populated by /insights/recompute-communities. Communities are ordered by size descending; one row per (community, member) pair so the caller can group as needed.", + responses = { + @ApiResponse(responseCode = "200", description = "Communities + members"), + @ApiResponse(responseCode = "400", description = "Invalid entityType or graphType"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response listCommunities( + @Context SecurityContext securityContext, + @Parameter(description = "Entity type whose community partition you want.", required = true) + @QueryParam("entityType") + @NotEmpty + String entityType, + @Parameter(description = "Source graph: lineage (default) or tagCoOccurrence.") + @QueryParam("graphType") + @DefaultValue("lineage") + String graphType) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + String sparql; + try { + sparql = CommunityComputation.listingSparql(entityType, graphType); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .type(MediaType.APPLICATION_JSON) + .build(); + } + return executeSparqlQuery(sparql, "json", "none"); + } + + @GET + @Path("/insights/path") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "findLineagePath", + summary = "Find the shortest lineage path between two entities", + description = + "BFS over the lineage graph (prov:wasDerivedFrom, om:upstream, om:downstream) returning the shortest path between two URIs. Use direction=upstream to walk from entity to its sources, downstream to walk to derived entities, both for either. Each hop returns the URI, the predicate that connected it, and any om:* rdf:type values. Useful for explain-lineage UIs and impact-analysis tooling.", + responses = { + @ApiResponse( + responseCode = "200", + description = "Path between the two entities (or found=false)"), + @ApiResponse( + responseCode = "400", + description = "Invalid from/to URI, direction, or maxHops"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response findLineagePath( + @Context SecurityContext securityContext, + @Parameter(description = "Starting entity URI (absolute http(s)).", required = true) + @QueryParam("from") + @NotEmpty + String from, + @Parameter(description = "Target entity URI (absolute http(s)).", required = true) + @QueryParam("to") + @NotEmpty + String to, + @Parameter(description = "Walk direction: upstream (default), downstream, or both.") + @QueryParam("direction") + @DefaultValue("upstream") + String direction, + @Parameter(description = "Max hops to explore. 1–25, defaults to 6.") @QueryParam("maxHops") + Integer maxHops) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + try { + LineagePathBuilder.Direction dir = LineagePathBuilder.Direction.parse(direction); + LineagePathFinder.Path path = + new LineagePathFinder(getRdfRepository()).findPath(from, to, dir, maxHops); + return Response.ok(JsonUtils.pojoToJson(path)).build(); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .type(MediaType.APPLICATION_JSON) + .build(); + } + } + + @GET + @Path("/insights/recommendations") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "datasetRecommendations", + summary = "Recommend related entities for a given seed URI", + description = + "Phase 3.4: ranks every other entity by graph-topology similarity to the given seed — overlap on tags, glossary terms, and direct lineage neighbours. Pure SPARQL, no precomputation. Score formula: 1.0 · tagOverlap + 1.5 · glossaryOverlap + 2.0 · lineageOverlap.", + responses = { + @ApiResponse(responseCode = "200", description = "Ranked recommendations"), + @ApiResponse(responseCode = "400", description = "Invalid entityUri or limit"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response datasetRecommendations( + @Context SecurityContext securityContext, + @Parameter(description = "Seed entity URI (absolute http(s)).", required = true) + @QueryParam("entityUri") + @NotEmpty + String entityUri, + @Parameter(description = "Number of recommendations. 1–50, default 10.") + @QueryParam("limit") + @DefaultValue("10") + int limit) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + String sparql; + try { + sparql = RecommendationsQueryBuilder.build(entityUri, limit); + } catch (IllegalArgumentException e) { + return Response.status(Response.Status.BAD_REQUEST) + .entity(buildErrorResponse(e.getMessage())) + .type(MediaType.APPLICATION_JSON) + .build(); + } + return executeSparqlQuery(sparql, "json", "none"); + } + + @GET + @Path("/insights/tag-cooccurrence") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "tagCoOccurrence", + summary = "Pairs of tags applied to the same entities", + description = + "Phase 3.5: returns pairs of tags that appear together on the same entity, sorted by overlap count descending. Surfaces governance signals like 'PII and Confidential are almost always co-applied'. Pure SPARQL aggregate over om:hasTag — no precomputation required.", + responses = { + @ApiResponse(responseCode = "200", description = "Tag pair counts"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response tagCoOccurrence( + @Context SecurityContext securityContext, + @Parameter(description = "Minimum number of shared entities. 1+, default 2.") + @QueryParam("minCount") + @DefaultValue("2") + int minCount, + @Parameter(description = "Number of pairs to return. 1–100, default 20.") + @QueryParam("limit") + @DefaultValue("20") + int limit) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + return executeSparqlQuery( + CoOccurrenceQueryBuilder.tagCoOccurrence(minCount, limit), "json", "none"); + } + + @GET + @Path("/insights/glossary-reach") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "glossaryReach", + summary = "Glossary terms ranked by domain reach", + description = + "Phase 3.5: returns glossary terms ordered by the number of distinct domains in which they're used, surfacing the most cross-cutting concepts. Pure SPARQL aggregate over om:hasGlossaryTerm × om:hasDomain.", + responses = { + @ApiResponse(responseCode = "200", description = "Term reach counts"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response glossaryReach( + @Context SecurityContext securityContext, + @Parameter(description = "Minimum number of domains. 1+, default 2.") + @QueryParam("minDomains") + @DefaultValue("2") + int minDomains, + @Parameter(description = "Number of terms to return. 1–100, default 20.") + @QueryParam("limit") + @DefaultValue("20") + int limit) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + return executeSparqlQuery( + CoOccurrenceQueryBuilder.glossaryReach(minDomains, limit), "json", "none"); + } + + @GET + @Path("/insights/tag-popularity") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "tagPopularity", + summary = "Tags ranked by number of tagged entities", + description = + "Phase 3.5: returns tags ordered by the number of distinct entities they're applied to. Companion to /insights/tag-cooccurrence — useful for triaging tag taxonomy bloat.", + responses = { + @ApiResponse(responseCode = "200", description = "Tag entity counts"), + @ApiResponse(responseCode = "503", description = "RDF service not enabled") + }) + public Response tagPopularity( + @Context SecurityContext securityContext, + @Parameter(description = "Number of tags to return. 1–100, default 20.") + @QueryParam("limit") + @DefaultValue("20") + int limit) { + authorizer.authorizeAdmin(securityContext); + if (getRdfRepository() == null || !getRdfRepository().isEnabled()) { + return Response.status(Response.Status.SERVICE_UNAVAILABLE) + .entity(buildErrorResponse("RDF repository is not enabled")) + .type(MediaType.APPLICATION_JSON) + .build(); + } + return executeSparqlQuery(CoOccurrenceQueryBuilder.tagPopularity(limit), "json", "none"); + } + + @GET + @Path("/ontology/extensions") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "listCustomOntologyExtensions", + summary = "List user-authored ontology extensions", + description = + "Returns every ontology extension registered with this server. Each extension is a bundle of custom OWL classes and properties under the om-extension namespace.", + responses = {@ApiResponse(responseCode = "200", description = "Extension list")}) + public Response listCustomOntologyExtensions(@Context SecurityContext securityContext) { + authorizer.authorizeAdmin(securityContext); + return Response.ok( + JsonUtils.pojoToJson(Map.of("extensions", CustomOntologyRegistry.getInstance().list()))) + .build(); + } + + @GET + @Path("/ontology/extensions/{name}") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "getCustomOntologyExtension", + summary = "Get a single custom ontology extension by name", + responses = { + @ApiResponse(responseCode = "200", description = "The extension"), + @ApiResponse(responseCode = "404", description = "Extension not found") + }) + public Response getCustomOntologyExtension( + @Context SecurityContext securityContext, @PathParam("name") String name) { + authorizer.authorizeAdmin(securityContext); + return CustomOntologyRegistry.getInstance() + .get(name) + .map(ext -> Response.ok(JsonUtils.pojoToJson(ext)).build()) + .orElse( + Response.status(Response.Status.NOT_FOUND) + .entity(buildErrorResponse("Custom ontology extension not found: " + name)) + .type(MediaType.APPLICATION_JSON) + .build()); + } + + @POST + @Path("/ontology/extensions/validate") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "validateCustomOntologyExtension", + summary = "Validate a candidate ontology extension without persisting it", + description = + "Runs the same validator that admin writes are gated on (URIs in om-extension namespace, no redefinition of canonical classes, no cycles, valid domain/range references) and returns the list of errors.", + responses = {@ApiResponse(responseCode = "200", description = "Validation result")}) + public Response validateCustomOntologyExtension( + @Context SecurityContext securityContext, CustomOntology candidate) { + authorizer.authorizeAdmin(securityContext); + List errors = CustomOntologyValidator.validate(candidate); + Map body = new java.util.LinkedHashMap<>(); + body.put("valid", errors.isEmpty()); + body.put("errors", errors); + return Response.ok(JsonUtils.pojoToJson(body)).build(); + } + + @POST + @Path("/rules/validate") + @Produces(MediaType.APPLICATION_JSON) + @Operation( + operationId = "validateInferenceRule", + summary = "Validate a candidate inference rule without persisting it", + description = + "Runs the same validator that admin writes are gated on (CONSTRUCT-only, no SERVICE clauses, syntactically well-formed) and returns the list of errors. Useful for an admin UI that wants live feedback while editing a rule.", + responses = { + @ApiResponse(responseCode = "200", description = "Validation result"), + @ApiResponse(responseCode = "403", description = "Forbidden") + }) + public Response validateInferenceRule( + @Context SecurityContext securityContext, InferenceRule candidate) { + authorizer.authorizeAdmin(securityContext); + List errors = InferenceRuleValidator.validate(candidate); + Map body = new java.util.LinkedHashMap<>(); + body.put("valid", errors.isEmpty()); + body.put("errors", errors); + return Response.ok(JsonUtils.pojoToJson(body)).build(); + } + + @GET + @Path("/ontology") + @Produces({TURTLE, RDF_XML, N_TRIPLES, JSON_LD, MediaType.WILDCARD}) + @Operation( + operationId = "getOntology", + summary = "Download the OpenMetadata ontology", + description = + "Returns the canonical OpenMetadata OWL ontology and its PROV-aligned extension as a single document. The ontology imports DCAT, PROV-O, and SKOS by reference. Format is selected via the Accept header or the format query param: turtle (default), rdfxml, ntriples, jsonld.", + responses = { + @ApiResponse( + responseCode = "200", + description = "Ontology document in the requested serialization", + content = { + @Content(mediaType = TURTLE), + @Content(mediaType = RDF_XML), + @Content(mediaType = N_TRIPLES), + @Content(mediaType = JSON_LD) + }), + @ApiResponse(responseCode = "500", description = "Ontology resource missing or unreadable") + }) + public Response getOntology( + @Parameter( + description = + "Output serialization. One of: turtle, rdfxml, ntriples, jsonld. Defaults to turtle.") + @QueryParam("format") + @DefaultValue("turtle") + String format) { + return OntologyDocument.serve(format); + } + @GET @Path("/debug/glossary-relations") @Operation( @@ -501,6 +1128,16 @@ private Response executeSparqlQuery(String query, String format, String inferenc .build(); } + try { + getFederationGuard().enforce(query); + } catch (SparqlFederationGuard.FederationDisallowedException e) { + LOG.warn("Rejected SPARQL with disallowed SERVICE clause: {}", e.getBlockedEndpoint()); + return Response.status(Response.Status.FORBIDDEN) + .entity(buildErrorResponse(e.getMessage())) + .type(MediaType.APPLICATION_JSON) + .build(); + } + try { String mimeType = getMimeTypeForFormat(format); String results; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfShaclValidator.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfShaclValidator.java new file mode 100644 index 000000000000..19338e6fa3f4 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/rdf/RdfShaclValidator.java @@ -0,0 +1,63 @@ +package org.openmetadata.service.resources.rdf; + +import java.io.IOException; +import java.io.InputStream; +import lombok.extern.slf4j.Slf4j; +import org.apache.jena.graph.Graph; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.ModelFactory; +import org.apache.jena.riot.Lang; +import org.apache.jena.riot.RDFDataMgr; +import org.apache.jena.shacl.Shapes; +import org.apache.jena.shacl.ValidationReport; + +/** + * Loads {@code openmetadata-shapes.ttl} from the classpath and validates an arbitrary RDF model + * against it. The validation is "report only" — callers decide whether a non-empty report is a + * blocker. The shapes graph is parsed once per JVM (Holder pattern). + */ +@Slf4j +public final class RdfShaclValidator { + + private static final String SHAPES_RESOURCE = "/rdf/shapes/openmetadata-shapes.ttl"; + + private RdfShaclValidator() {} + + private static final class Holder { + private static final Shapes SHAPES = loadShapes(); + } + + private static Shapes loadShapes() { + Model shapesModel = ModelFactory.createDefaultModel(); + try (InputStream is = RdfShaclValidator.class.getResourceAsStream(SHAPES_RESOURCE)) { + if (is == null) { + LOG.warn("SHACL shapes resource not found on classpath: {}", SHAPES_RESOURCE); + return Shapes.parse(shapesModel.getGraph()); + } + RDFDataMgr.read(shapesModel, is, Lang.TURTLE); + } catch (IOException e) { + LOG.warn("Failed to read SHACL shapes resource {}: {}", SHAPES_RESOURCE, e.getMessage()); + } catch (RuntimeException e) { + // RDFDataMgr.read can throw RiotException (and other Jena RuntimeExceptions) on a malformed + // TTL. Catch broadly so a corrupt resource degrades to an empty shape set rather than + // failing class initialization and taking down callers (RdfResource, MCP tools). + LOG.warn("Failed to parse SHACL shapes resource {}: {}", SHAPES_RESOURCE, e.getMessage()); + shapesModel = ModelFactory.createDefaultModel(); + } + return Shapes.parse(shapesModel.getGraph()); + } + + /** Validate {@code data} against the OpenMetadata shapes. Never throws on conforming data. */ + public static ValidationReport validate(Graph data) { + return org.apache.jena.shacl.ShaclValidator.get().validate(Holder.SHAPES, data); + } + + public static ValidationReport validate(Model data) { + return validate(data.getGraph()); + } + + /** Expose the shapes for callers that need to inspect or extend them. */ + public static Shapes shapes() { + return Holder.SHAPES; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataOperations.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataOperations.java index 52701b14bd90..66a8d79c2263 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataOperations.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataOperations.java @@ -133,16 +133,6 @@ import org.openmetadata.service.secrets.SecretsManagerUpdateService; import org.openmetadata.service.security.auth.SecurityConfigurationManager; import org.openmetadata.service.security.jwt.JWTTokenGenerator; -import org.openmetadata.service.util.dbtune.AutoTuner; -import org.openmetadata.service.util.dbtune.DbTuneDiagnosis; -import org.openmetadata.service.util.dbtune.DbTuneReport; -import org.openmetadata.service.util.dbtune.DbTuneResult; -import org.openmetadata.service.util.dbtune.Diagnostic; -import org.openmetadata.service.util.dbtune.MysqlAutoTuner; -import org.openmetadata.service.util.dbtune.MysqlDiagnostic; -import org.openmetadata.service.util.dbtune.PostgresAutoTuner; -import org.openmetadata.service.util.dbtune.PostgresDiagnostic; -import org.openmetadata.service.util.dbtune.TableRecommendation; import org.openmetadata.service.util.jdbi.DatabaseAuthenticationProviderFactory; import org.openmetadata.service.util.jdbi.JdbiUtils; import org.slf4j.LoggerFactory; @@ -185,13 +175,9 @@ public Integer call() { + "'drop-create', 'changelog', 'migrate', 'migrate-secrets', 'reindex', 'reembed', 'reindex-rdf', 'reindexdi', 'deploy-pipelines', " + "'dbServiceCleanup', 'relationshipCleanup', 'tagUsageCleanup', 'drop-indexes', 'remove-security-config', 'create-indexes', " + "'setOpenMetadataUrl', 'configureEmailSettings', 'get-security-config', 'update-security-config', 'install-app', 'delete-app', 'create-user', 'reset-password', " - + "'syncAlertOffset', 'analyze-tables', 'db-tune', 'cleanup-flowable-history', 'regenerate-bot-tokens'"); + + "'syncAlertOffset', 'analyze-tables', 'cleanup-flowable-history', 'regenerate-bot-tokens'"); LOG.info( "Use 'reindex --auto-tune' for automatic performance optimization based on cluster capabilities"); - LOG.info( - "Use 'db-tune' for a per-table autovacuum / InnoDB stats tuning report; add --apply to " - + "execute the recommendations, --analyze to refresh planner stats on changed tables, " - + "and --diagnose to surface unused indexes, bloat, slow queries, and other DBA findings"); LOG.info( "Use 'cleanup-flowable-history --delete --runtime-batch-size=1000 --history-batch-size=1000' for Flowable cleanup with custom options"); LOG.info( @@ -2483,134 +2469,6 @@ public Integer analyzeTables() { } } - @Command( - name = "db-tune", - description = - "Generate a per-table autovacuum / InnoDB stats tuning report and optionally apply it. " - + "Default mode is read-only — pass --apply to execute the ALTER TABLE statements, " - + "--analyze to refresh planner stats on changed tables, and --diagnose to also " - + "surface unused indexes, bloat, slow queries, and other read-only DBA findings.") - public Integer dbTune( - @Option( - names = {"--apply"}, - defaultValue = "false", - description = - "Apply the recommendations. Without this flag the command only prints the report.") - boolean apply, - @Option( - names = {"--yes", "-y"}, - defaultValue = "false", - description = "Skip the interactive confirmation when applying.") - boolean skipPrompt, - @Option( - names = {"--analyze"}, - defaultValue = "false", - description = - "After --apply, run ANALYZE on each changed table so planner stats reflect the new settings.") - boolean runAnalyze, - @Option( - names = {"--diagnose"}, - defaultValue = "false", - description = - "Also run a read-only diagnostic pass (unused indexes, bloat, low cache hit, " - + "stale ANALYZE, seq-scan-heavy tables, slow queries). Pure inspection — " - + "never modifies anything.") - boolean runDiagnose) { - try { - parseConfig(); - String driverClass = config.getDataSourceFactory().getDriverClass(); - ConnectionType connType = ConnectionType.from(driverClass); - if (connType == null) { - LOG.error( - "db-tune does not support driver class '{}'. Only the bundled MySQL and PostgreSQL drivers are recognised.", - driverClass); - return 1; - } - AutoTuner tuner = autoTunerFor(connType); - DbTuneResult result = jdbi.withHandle(tuner::analyze); - LOG.info("\n{}", DbTuneReport.render(result)); - if (runDiagnose) { - Diagnostic diagnostic = diagnosticFor(connType); - DbTuneDiagnosis diagnosis = jdbi.withHandle(diagnostic::diagnose); - LOG.info("\n{}", DbTuneReport.renderDiagnosis(diagnosis)); - } - if (!apply) { - return 0; - } - List actionable = result.actionableRecommendations(); - if (actionable.isEmpty()) { - if (result.tableRecommendations().isEmpty()) { - LOG.info("Nothing to apply — no tracked tables exist on this database."); - } else { - LOG.info( - "Nothing to apply — every tracked table already matches its recommended settings."); - } - return 0; - } - if (!skipPrompt && !confirmApply(tuner, actionable)) { - LOG.info("Operation cancelled."); - return 0; - } - applyRecommendations(tuner, actionable, runAnalyze); - return 0; - } catch (Exception e) { - LOG.error("db-tune failed due to ", e); - return 1; - } - } - - private AutoTuner autoTunerFor(final ConnectionType connType) { - return switch (connType) { - case POSTGRES -> new PostgresAutoTuner(); - case MYSQL -> new MysqlAutoTuner(); - }; - } - - private Diagnostic diagnosticFor(final ConnectionType connType) { - return switch (connType) { - case POSTGRES -> new PostgresDiagnostic(); - case MYSQL -> new MysqlDiagnostic(); - }; - } - - private boolean confirmApply(final AutoTuner tuner, final List actionable) { - LOG.info("About to apply {} ALTER statements:", actionable.size()); - LOG.info("\n{}", DbTuneReport.renderAlterStatements(tuner, actionable)); - @SuppressWarnings("resource") - Scanner scanner = new Scanner(System.in); - LOG.info("Apply now? [y/N]: "); - // nextLine() (not next()) so a bare Enter — which the [y/N] convention implies as "no" — - // doesn't block waiting for a non-whitespace token. Treat empty / EOF as "no". - String input = scanner.hasNextLine() ? scanner.nextLine().trim().toLowerCase() : ""; - return input.equals("y") || input.equals("yes"); - } - - private void applyRecommendations( - final AutoTuner tuner, final List actionable, final boolean runAnalyze) { - List> rows = new ArrayList<>(); - for (TableRecommendation rec : actionable) { - rows.add(applyOne(tuner, rec, runAnalyze)); - } - printToAsciiTable( - List.of("Table", "Action", "Status", "Details"), rows, "No recommendations applied"); - } - - private List applyOne( - final AutoTuner tuner, final TableRecommendation rec, final boolean runAnalyze) { - try { - jdbi.useHandle(handle -> tuner.apply(handle, rec)); - if (runAnalyze) { - jdbi.useHandle(handle -> tuner.analyzeOne(handle, rec.tableName())); - return List.of(rec.tableName(), rec.action().name(), "OK", "Applied + analyzed"); - } - return List.of(rec.tableName(), rec.action().name(), "OK", "Applied"); - } catch (Exception e) { - LOG.error("Failed to apply recommendation for {}: {}", rec.tableName(), e.getMessage(), e); - String detail = e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName(); - return List.of(rec.tableName(), rec.action().name(), "FAILED", detail); - } - } - /** * Unlike most ops commands (e.g. deploy-pipelines) that delegate to the server API, this command * operates directly on the database. This is intentional: when JWT signing keys have been rotated, diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Action.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Action.java deleted file mode 100644 index 199115fe65d8..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Action.java +++ /dev/null @@ -1,25 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -public enum Action { - APPLY, - TIGHTEN, - RELAX, - OK, - SKIP; - - public boolean isActionable() { - return this == APPLY || this == TIGHTEN || this == RELAX; - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/AutoTuner.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/AutoTuner.java deleted file mode 100644 index f03c568b71ad..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/AutoTuner.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.Map; -import org.jdbi.v3.core.Handle; - -/** - * Engine-specific auto-tuner. Implementations: - * - *
    - *
  • Read observed table stats and current parameter-group settings from the database. - *
  • Compute a recommended table-level reloption set per table (pure logic — see - * {@link #recommend(TableStats)}). - *
  • Apply the recommendations via {@code ALTER TABLE ... SET (...)} when the operator opts in. - *
  • Optionally refresh planner stats on tables that were changed. - *
- */ -public interface AutoTuner { - - /** Read stats + settings, then turn them into recommendations. Mixes I/O and pure logic. */ - DbTuneResult analyze(Handle handle); - - /** - * Reads the current per-table reloption / table-option settings for an arbitrary table name — - * including tables that are not in the static tuning catalog. Returns an empty map if the table - * has no overrides set (inherits cluster defaults). The returned keys use the same casing the - * engine reports (lowercase autovacuum_* keys for Postgres, uppercase STATS_* keys for MySQL). - */ - Map currentSettingsForTable(Handle handle, String tableName); - - /** - * Pure decision function. Given observed table stats, return the recommendation. Exposed - * separately so unit tests can assert the heuristic without hitting a database. - */ - TableRecommendation recommend(TableStats stats); - - /** - * Apply a single actionable recommendation. No-op for non-actionable actions. Idempotent — safe - * to re-run. - */ - void apply(Handle handle, TableRecommendation recommendation); - - /** Refresh planner stats for one table after a settings change. */ - void analyzeOne(Handle handle, String tableName); - - /** Build the {@code ALTER TABLE} statement for a recommendation. Engine-specific syntax. */ - String buildAlterStatement(TableRecommendation recommendation); -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneDiagnosis.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneDiagnosis.java deleted file mode 100644 index a2e88ac006a9..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneDiagnosis.java +++ /dev/null @@ -1,36 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; - -/** Diagnostic result bundle. {@code notes} carries advisory messages (e.g. missing extension). */ -public record DbTuneDiagnosis(List findings, List notes) { - - public DbTuneDiagnosis { - findings = findings == null ? List.of() : List.copyOf(findings); - notes = notes == null ? List.of() : List.copyOf(notes); - } - - /** Group findings by category preserving the enum order so the report sections print stably. */ - public Map> findingsByCategory() { - return findings.stream() - .collect( - Collectors.groupingBy( - Finding::category, - () -> new java.util.EnumMap<>(DiagnosticCategory.class), - Collectors.toList())); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneReport.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneReport.java deleted file mode 100644 index 439e51959c3a..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneReport.java +++ /dev/null @@ -1,203 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.text.NumberFormat; -import java.util.ArrayList; -import java.util.List; -import java.util.Locale; -import java.util.Map; -import java.util.stream.Collectors; -import org.openmetadata.service.util.AsciiTable; - -public final class DbTuneReport { - - private static final NumberFormat ROW_FORMAT = NumberFormat.getInstance(Locale.ROOT); - private static final long KB = 1024L; - private static final long MB = KB * 1024L; - private static final long GB = MB * 1024L; - - private DbTuneReport() {} - - public static String render(final DbTuneResult result) { - StringBuilder out = new StringBuilder(); - out.append("Database engine: ").append(result.engine()); - if (result.engineVersion() != null && !result.engineVersion().isBlank()) { - out.append(" ").append(result.engineVersion()); - } - out.append('\n').append('\n'); - appendServerParams(out, result.serverParams()); - appendTableRecommendations(out, result.tableRecommendations()); - appendNextSteps( - out, result.tableRecommendations().size(), result.actionableRecommendations().size()); - return out.toString(); - } - - private static void appendServerParams( - final StringBuilder out, final List checks) { - out.append("=== Server-level parameter compliance ===\n"); - if (checks.isEmpty()) { - out.append("(no parameter-group checks for this engine)\n\n"); - return; - } - List headers = List.of("Parameter", "Current", "Recommended", "Status", "Note"); - List> rows = - checks.stream() - .map( - c -> - List.of( - nullToBlank(c.parameter()), - nullToBlank(c.currentValue()), - nullToBlank(c.recommendedValue()), - nullToBlank(c.status()), - nullToBlank(c.note()))) - .toList(); - out.append(new AsciiTable(headers, rows, true, "", "(empty)").render()); - out.append('\n'); - out.append( - "These cannot be applied by this tool — change them in your DB parameter group / RDS console.\n\n"); - } - - private static void appendTableRecommendations( - final StringBuilder out, final List recs) { - out.append("=== Per-table recommendations (").append(recs.size()).append(" tables) ===\n"); - if (recs.isEmpty()) { - out.append("(no recommendations — none of the tracked tables exist on this database)\n\n"); - return; - } - List headers = - List.of("Table", "Rows", "Size", "Current", "Recommended", "Action", "Reason"); - List> rows = - recs.stream() - .map( - r -> - List.of( - r.tableName(), - ROW_FORMAT.format(r.rowCount()), - formatBytes(r.totalBytes()), - formatSettings(r.currentSettings()), - formatSettings(r.recommendedSettings()), - r.action().name(), - nullToBlank(r.reason()))) - .toList(); - out.append(new AsciiTable(headers, rows, true, "", "(empty)").render()); - out.append('\n'); - } - - private static void appendNextSteps( - final StringBuilder out, final int totalRecommendations, final int actionableCount) { - if (totalRecommendations == 0) { - // No tracked tables exist on this database — saying "all match" would be misleading. - return; - } - if (actionableCount == 0) { - out.append("All tracked tables already match their recommended settings — nothing to do.\n"); - return; - } - out.append("Next steps:\n"); - out.append( - " ./bootstrap/openmetadata-ops.sh db-tune --apply --analyze # apply + refresh planner stats\n"); - out.append( - " ./bootstrap/openmetadata-ops.sh db-tune --apply # apply only; run analyze-tables later\n"); - } - - static String formatSettings(final Map settings) { - if (settings == null || settings.isEmpty()) { - return "(default)"; - } - return settings.entrySet().stream() - .sorted(Map.Entry.comparingByKey()) - .map(e -> e.getKey() + "=" + e.getValue()) - .collect(Collectors.joining(", ")); - } - - static String formatBytes(final long bytes) { - if (bytes <= 0) { - return "0 B"; - } - if (bytes >= GB) { - return String.format(Locale.ROOT, "%.1f GB", bytes / (double) GB); - } - if (bytes >= MB) { - return String.format(Locale.ROOT, "%.0f MB", bytes / (double) MB); - } - if (bytes >= KB) { - return String.format(Locale.ROOT, "%.0f KB", bytes / (double) KB); - } - return bytes + " B"; - } - - private static String nullToBlank(final String value) { - return value == null ? "" : value; - } - - /** Concatenates each recommendation's ALTER statement, one per line, terminated by a semicolon. */ - public static String renderAlterStatements( - final AutoTuner tuner, final List recommendations) { - List lines = new ArrayList<>(recommendations.size()); - for (TableRecommendation rec : recommendations) { - lines.add(tuner.buildAlterStatement(rec) + ";"); - } - return String.join("\n", lines); - } - - /** - * Renders read-only diagnostic findings grouped by category. Each category that produced at - * least one finding gets its own section with a category-specific column layout. Categories with - * zero findings are suppressed; the {@code notes} list is appended at the end so an operator sees - * what couldn't be checked (missing extension, permissions, etc.). - */ - public static String renderDiagnosis(final DbTuneDiagnosis diagnosis) { - StringBuilder out = new StringBuilder(); - out.append("=== Diagnostic findings ===\n"); - Map> grouped = diagnosis.findingsByCategory(); - if (grouped.isEmpty()) { - out.append("(no findings — every check returned a clean result)\n"); - } - for (Map.Entry> e : grouped.entrySet()) { - appendCategorySection(out, e.getKey(), e.getValue()); - } - appendNotes(out, diagnosis.notes()); - return out.toString(); - } - - private static void appendCategorySection( - final StringBuilder out, final DiagnosticCategory category, final List findings) { - out.append('\n') - .append(category.title()) - .append(" (") - .append(findings.size()) - .append(" found):\n"); - out.append(" ").append(category.description()).append('\n'); - List> rows = new ArrayList<>(); - for (Finding f : findings) { - List row = new ArrayList<>(category.columns().size()); - for (String col : category.columns()) { - row.add(nullToBlank(f.attributes().get(col))); - } - rows.add(row); - } - out.append(new AsciiTable(category.columns(), rows, true, "", "(empty)").render()); - out.append('\n'); - } - - private static void appendNotes(final StringBuilder out, final List notes) { - if (notes == null || notes.isEmpty()) { - return; - } - out.append("\nNotes:\n"); - for (String note : notes) { - out.append(" - ").append(note).append('\n'); - } - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneResult.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneResult.java deleted file mode 100644 index ead95f324ae8..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DbTuneResult.java +++ /dev/null @@ -1,32 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.List; - -public record DbTuneResult( - String engine, - String engineVersion, - List serverParams, - List tableRecommendations) { - - public DbTuneResult { - serverParams = serverParams == null ? List.of() : List.copyOf(serverParams); - tableRecommendations = - tableRecommendations == null ? List.of() : List.copyOf(tableRecommendations); - } - - public List actionableRecommendations() { - return tableRecommendations.stream().filter(r -> r.action().isActionable()).toList(); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Diagnostic.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Diagnostic.java deleted file mode 100644 index 9e10f2af851d..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Diagnostic.java +++ /dev/null @@ -1,26 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import org.jdbi.v3.core.Handle; - -/** - * Read-only DBA diagnostic. Inspects the live database for unused indexes, bloat indicators, slow - * queries, and other signals. Implementations must catch and log per-category errors so a missing - * extension (e.g. {@code pg_stat_statements} not installed) does not abort the whole diagnose run - * — surface it in {@link DbTuneDiagnosis#notes()} instead. - */ -public interface Diagnostic { - - DbTuneDiagnosis diagnose(Handle handle); -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DiagnosticCategory.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DiagnosticCategory.java deleted file mode 100644 index c7bc9d2621f4..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/DiagnosticCategory.java +++ /dev/null @@ -1,77 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.List; - -/** - * Categories of read-only diagnostic findings emitted by {@link Diagnostic#diagnose}. Each category - * has a fixed list of attribute keys that {@link Finding#attributes} is expected to populate; the - * report renderer dispatches column layout per category. - */ -public enum DiagnosticCategory { - UNUSED_INDEX( - "Unused indexes", - "Indexes with zero scans since last stats reset; candidates for DROP after a usage review.", - List.of("table", "index", "size", "scans")), - HIGH_DEAD_TUPLES( - "Tables with high dead-tuple ratio", - "n_dead_tup / n_live_tup > 0.2 — autovacuum is falling behind on this table.", - List.of("table", "live_rows", "dead_rows", "dead_ratio", "last_vacuum")), - LOW_CACHE_HIT( - "Tables with low cache hit ratio", - "Heap reads exceed 1000 with hit ratio < 90%; suggests undersized buffers or hot seq scans.", - List.of("table", "heap_reads", "heap_hits", "hit_pct")), - STALE_STATS( - "Tables with stale ANALYZE", - "Last autoanalyze older than 14 days (or never); planner stats may be misleading.", - List.of("table", "last_analyzed", "live_rows")), - SEQ_SCAN_HEAVY( - "Tables with seq-scan-heavy access", - "seq_scan/idx_scan > 10 with > 1000 seq scans; suggests a missing index.", - List.of("table", "seq_scans", "idx_scans", "ratio")), - SLOW_QUERY( - "Top slowest queries", - "From pg_stat_statements / events_statements_summary_by_digest. Truncated to 100 chars.", - List.of("query", "calls", "mean_ms")), - FULL_TABLE_SCAN( - "Queries doing full table scans", - "From sys.statements_with_full_table_scans (MySQL).", - List.of("query", "exec_count", "rows_examined_avg")), - LOW_BUFFER_POOL_HIT( - "InnoDB buffer pool hit ratio", - "Hit ratio < 99% suggests undersized innodb_buffer_pool_size for the working set.", - List.of("metric", "value")); - - private final String title; - private final String description; - private final List columns; - - DiagnosticCategory(final String title, final String description, final List columns) { - this.title = title; - this.description = description; - this.columns = List.copyOf(columns); - } - - public String title() { - return title; - } - - public String description() { - return description; - } - - public List columns() { - return columns; - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Finding.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Finding.java deleted file mode 100644 index 7026f9df5430..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Finding.java +++ /dev/null @@ -1,28 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.Map; - -/** - * One row of a diagnostic finding. {@code attributes} keys must match {@link - * DiagnosticCategory#columns()} for the same {@code category} so the renderer can lay them out - * predictably. - */ -public record Finding( - DiagnosticCategory category, Severity severity, Map attributes) { - - public Finding { - attributes = attributes == null ? Map.of() : Map.copyOf(attributes); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlAutoTuner.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlAutoTuner.java deleted file mode 100644 index 5c8102f26701..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlAutoTuner.java +++ /dev/null @@ -1,266 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.ArrayList; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Locale; -import java.util.Map; -import java.util.stream.Collectors; -import org.jdbi.v3.core.Handle; -import org.openmetadata.service.util.dbtune.MysqlTuningCatalog.Profile; - -public final class MysqlAutoTuner implements AutoTuner { - - @Override - public DbTuneResult analyze(final Handle handle) { - String version = readVersion(handle); - List serverParams = readServerParams(handle); - List stats = loadTableStats(handle); - List recs = stats.stream().map(this::recommend).toList(); - return new DbTuneResult("MySQL", version, serverParams, recs); - } - - @Override - public TableRecommendation recommend(final TableStats stats) { - Profile profile = MysqlTuningCatalog.profileFor(stats.tableName()); - if (profile == null) { - return skip(stats, "Table is not in the dbtune catalog"); - } - if (stats.rowCount() < profile.rowThreshold()) { - return skip( - stats, - String.format( - Locale.ROOT, - "Row count %d below threshold %d", - stats.rowCount(), - profile.rowThreshold())); - } - return decideAction(stats, profile); - } - - private TableRecommendation decideAction(final TableStats stats, final Profile profile) { - Map recommended = profile.settings(); - Map current = stats.currentSettings(); - if (settingsMatch(current, recommended)) { - return new TableRecommendation( - stats.tableName(), - Action.OK, - stats.rowCount(), - stats.totalBytes(), - current, - recommended, - "Already matches recommended settings"); - } - Action action = current.isEmpty() ? Action.APPLY : Action.TIGHTEN; - return new TableRecommendation( - stats.tableName(), - action, - stats.rowCount(), - stats.totalBytes(), - current, - recommended, - profile.reason()); - } - - @Override - public void apply(final Handle handle, final TableRecommendation recommendation) { - if (!recommendation.action().isActionable()) { - return; - } - handle.execute(buildAlterStatement(recommendation)); - } - - @Override - public void analyzeOne(final Handle handle, final String tableName) { - handle.execute("ANALYZE TABLE " + quoteIdent(tableName)); - } - - @Override - public String buildAlterStatement(final TableRecommendation recommendation) { - String settings = - recommendation.recommendedSettings().entrySet().stream() - .sorted(Map.Entry.comparingByKey()) - .map(e -> e.getKey() + "=" + e.getValue()) - .collect(Collectors.joining(", ")); - return "ALTER TABLE " + quoteIdent(recommendation.tableName()) + " " + settings; - } - - // ---- DB I/O ---- - - String readVersion(final Handle handle) { - return handle.createQuery("SELECT VERSION()").mapTo(String.class).findOne().orElse(""); - } - - List loadTableStats(final Handle handle) { - List result = new ArrayList<>(); - for (String tableName : MysqlTuningCatalog.tableNames()) { - TableStats stats = loadTableStats(handle, tableName); - if (stats != null) { - result.add(stats); - } - } - return result; - } - - TableStats loadTableStats(final Handle handle, final String tableName) { - return handle - .createQuery( - "SELECT TABLE_ROWS AS rows_estimate, " - + " COALESCE(DATA_LENGTH, 0) AS heap_bytes, " - + " COALESCE(INDEX_LENGTH, 0) AS idx_bytes, " - + " COALESCE(CREATE_OPTIONS, '') AS create_opts " - + "FROM information_schema.TABLES " - + "WHERE TABLE_SCHEMA = DATABASE() " - + " AND TABLE_NAME = :name") - .bind("name", tableName) - .map( - (rs, ctx) -> - new TableStats( - tableName, - Math.max(rs.getLong("rows_estimate"), 0), - rs.getLong("heap_bytes"), - rs.getLong("idx_bytes"), - parseCreateOptions(rs.getString("create_opts")))) - .findOne() - .orElse(null); - } - - @Override - public Map currentSettingsForTable(final Handle handle, final String tableName) { - return handle - .createQuery( - "SELECT COALESCE(CREATE_OPTIONS, '') AS create_opts " - + "FROM information_schema.TABLES " - + "WHERE TABLE_SCHEMA = DATABASE() " - + " AND TABLE_NAME = :name") - .bind("name", tableName) - .mapTo(String.class) - .findOne() - .map(MysqlAutoTuner::parseCreateOptions) - .orElse(Map.of()); - } - - List readServerParams(final Handle handle) { - List checks = new ArrayList<>(); - Map recommendations = recommendedServerParams(); - for (Map.Entry e : recommendations.entrySet()) { - String name = e.getKey(); - String recommended = e.getValue(); - String current = readGlobalVariable(handle, name); - checks.add(buildServerCheck(name, current, recommended)); - } - return checks; - } - - // ---- helpers ---- - - static Map parseCreateOptions(final String createOptions) { - if (createOptions == null || createOptions.isBlank()) { - return Map.of(); - } - Map out = new LinkedHashMap<>(); - for (String token : createOptions.trim().split("\\s+")) { - int eq = token.indexOf('='); - if (eq > 0) { - String key = token.substring(0, eq).toUpperCase(Locale.ROOT); - String value = token.substring(eq + 1); - if (key.startsWith("STATS_")) { - out.put(key, value); - } - } - } - return Map.copyOf(out); - } - - static boolean settingsMatch(final Map current, final Map rec) { - for (Map.Entry e : rec.entrySet()) { - String currentValue = current.get(e.getKey()); - if (currentValue == null || !numericEquals(currentValue, e.getValue())) { - return false; - } - } - return true; - } - - private static boolean numericEquals(final String a, final String b) { - try { - return Double.parseDouble(a) == Double.parseDouble(b); - } catch (NumberFormatException ex) { - return a.equalsIgnoreCase(b); - } - } - - private static TableRecommendation skip(final TableStats stats, final String reason) { - return new TableRecommendation( - stats.tableName(), - Action.SKIP, - stats.rowCount(), - stats.totalBytes(), - stats.currentSettings(), - Map.of(), - reason); - } - - static String quoteIdent(final String identifier) { - if (!identifier.matches("[a-zA-Z_][a-zA-Z0-9_]*")) { - throw new IllegalArgumentException( - "Refusing to build SQL with unsafe identifier: " + identifier); - } - return "`" + identifier + "`"; - } - - private String readGlobalVariable(final Handle handle, final String name) { - return handle - .createQuery( - "SELECT VARIABLE_VALUE FROM performance_schema.global_variables " - + "WHERE VARIABLE_NAME = :n") - .bind("n", name.toLowerCase(Locale.ROOT)) - .mapTo(String.class) - .findOne() - .orElse(null); - } - - static Map recommendedServerParams() { - Map map = new LinkedHashMap<>(); - map.put("innodb_buffer_pool_size", "40-60% of RAM (use formula form on RDS)"); - map.put("innodb_io_capacity", "2000"); - map.put("innodb_io_capacity_max", "4000"); - map.put("innodb_stats_persistent_sample_pages", "64"); - map.put("sort_buffer_size", "8388608"); // 8 MB - map.put("join_buffer_size", "4194304"); // 4 MB - map.put("tmp_table_size", "67108864"); // 64 MB - map.put("max_heap_table_size", "67108864"); // 64 MB - return Map.copyOf(map); - } - - static ServerParamCheck buildServerCheck( - final String name, final String current, final String recommended) { - if (current == null) { - return new ServerParamCheck( - name, "", recommended, ServerParamCheck.STATUS_UNKNOWN, "Variable not visible"); - } - if (recommended.contains("%")) { - return new ServerParamCheck( - name, - current, - recommended, - ServerParamCheck.STATUS_UNTUNED, - "RAM-relative; verify in RDS"); - } - boolean ok = numericEquals(current, recommended); - String status = ok ? ServerParamCheck.STATUS_OK : ServerParamCheck.STATUS_MISMATCH; - return new ServerParamCheck(name, current, recommended, status, ""); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlDiagnostic.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlDiagnostic.java deleted file mode 100644 index e957be7db2e4..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlDiagnostic.java +++ /dev/null @@ -1,189 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.ArrayList; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Locale; -import java.util.Map; -import lombok.extern.slf4j.Slf4j; -import org.jdbi.v3.core.Handle; - -/** - * MySQL diagnostic. Reads from {@code sys.*}, {@code performance_schema.*}, and - * {@code INFORMATION_SCHEMA} views; gracefully degrades if a view is missing or permissions are - * insufficient (the operator gets a {@link DbTuneDiagnosis#notes()} entry). - */ -@Slf4j -public final class MysqlDiagnostic implements Diagnostic { - - static final double LOW_BUFFER_POOL_HIT = 0.99; - static final int SLOW_QUERY_LIMIT = 10; - static final int QUERY_TRUNCATE = 100; - - @Override - public DbTuneDiagnosis diagnose(final Handle handle) { - List findings = new ArrayList<>(); - List notes = new ArrayList<>(); - runCategory(handle, notes, "unused indexes", h -> findings.addAll(unusedIndexes(h))); - runCategory(handle, notes, "buffer pool hit", h -> findings.addAll(bufferPoolHit(h, notes))); - runCategory(handle, notes, "slow queries", h -> findings.addAll(slowQueries(h, notes))); - runCategory(handle, notes, "full table scans", h -> findings.addAll(fullTableScans(h, notes))); - return new DbTuneDiagnosis(findings, notes); - } - - private void runCategory( - final Handle handle, - final List notes, - final String label, - final java.util.function.Consumer body) { - try { - body.accept(handle); - } catch (Exception e) { - LOG.warn("Diagnostic [{}] failed: {}", label, e.getMessage()); - notes.add(label + ": " + e.getMessage()); - } - } - - // ---- categories ---- - - List unusedIndexes(final Handle handle) { - return handle - .createQuery( - "SELECT object_schema, object_name, index_name " - + "FROM sys.schema_unused_indexes " - + "WHERE object_schema = DATABASE() " - + "ORDER BY object_name " - + "LIMIT 50") - .map( - (rs, ctx) -> - new Finding( - DiagnosticCategory.UNUSED_INDEX, - Severity.WARN, - Map.of( - "table", - rs.getString("object_name"), - "index", - rs.getString("index_name"), - "size", - "(not in view)", - "scans", - "0"))) - .list(); - } - - List bufferPoolHit(final Handle handle, final List notes) { - Long reads = readGlobalStatusLong(handle, "Innodb_buffer_pool_reads"); - Long requests = readGlobalStatusLong(handle, "Innodb_buffer_pool_read_requests"); - if (reads == null || requests == null || requests == 0) { - notes.add("buffer pool hit: Innodb_buffer_pool_* counters not available"); - return List.of(); - } - double hitRatio = 1.0 - (reads.doubleValue() / requests.doubleValue()); - if (hitRatio >= LOW_BUFFER_POOL_HIT) { - return List.of(); - } - return List.of( - new Finding( - DiagnosticCategory.LOW_BUFFER_POOL_HIT, - Severity.INFO, - Map.of( - "metric", - "innodb_buffer_pool_hit_ratio", - "value", - String.format(Locale.ROOT, "%.4f", hitRatio)))); - } - - List slowQueries(final Handle handle, final List notes) { - try { - return handle - .createQuery( - "SELECT digest_text, count_star AS calls, " - + " ROUND(avg_timer_wait/1000000, 2) AS mean_us " - + "FROM performance_schema.events_statements_summary_by_digest " - + "WHERE schema_name = DATABASE() " - + " AND digest_text IS NOT NULL " - + "ORDER BY avg_timer_wait DESC " - + "LIMIT :limit") - .bind("limit", SLOW_QUERY_LIMIT) - .map( - (rs, ctx) -> { - Map attrs = new LinkedHashMap<>(); - attrs.put("query", truncate(rs.getString("digest_text"))); - attrs.put("calls", String.valueOf(rs.getLong("calls"))); - attrs.put( - "mean_ms", - String.format(Locale.ROOT, "%.2f", rs.getDouble("mean_us") / 1000.0)); - return new Finding(DiagnosticCategory.SLOW_QUERY, Severity.INFO, attrs); - }) - .list(); - } catch (Exception e) { - notes.add("slow queries: performance_schema not available (" + e.getMessage() + ")"); - return List.of(); - } - } - - List fullTableScans(final Handle handle, final List notes) { - try { - return handle - .createQuery( - "SELECT query, exec_count, rows_examined_avg " - + "FROM sys.statements_with_full_table_scans " - + "WHERE db = DATABASE() " - + "ORDER BY exec_count DESC " - + "LIMIT 10") - .map( - (rs, ctx) -> - new Finding( - DiagnosticCategory.FULL_TABLE_SCAN, - Severity.INFO, - Map.of( - "query", truncate(rs.getString("query")), - "exec_count", String.valueOf(rs.getLong("exec_count")), - "rows_examined_avg", String.valueOf(rs.getLong("rows_examined_avg"))))) - .list(); - } catch (Exception e) { - notes.add( - "full table scans: sys.statements_with_full_table_scans not available (" - + e.getMessage() - + ")"); - return List.of(); - } - } - - private Long readGlobalStatusLong(final Handle handle, final String name) { - try { - return handle - .createQuery( - "SELECT VARIABLE_VALUE FROM performance_schema.global_status " - + "WHERE VARIABLE_NAME = :n") - .bind("n", name) - .mapTo(Long.class) - .findOne() - .orElse(null); - } catch (Exception e) { - return null; - } - } - - static String truncate(final String query) { - if (query == null) { - return ""; - } - String collapsed = query.replaceAll("\\s+", " ").trim(); - return collapsed.length() <= QUERY_TRUNCATE - ? collapsed - : collapsed.substring(0, QUERY_TRUNCATE) + "…"; - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlTuningCatalog.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlTuningCatalog.java deleted file mode 100644 index 8bac39f1e13c..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/MysqlTuningCatalog.java +++ /dev/null @@ -1,130 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.Set; - -/** - * Static catalog of which tables get which MySQL/InnoDB persistent-stats reloptions, and the - * row-count threshold below which we skip tuning that table. - * - *

InnoDB does not expose autovacuum knobs at the per-table level — purge is global. The lever - * that DOES help on large hot tables is bumping {@code STATS_SAMPLE_PAGES} above the default 20 so - * the planner picks the right index against multi-GB JSONB heaps. {@code STATS_PERSISTENT=1} + - * {@code STATS_AUTO_RECALC=1} are the modern InnoDB defaults; we assert them explicitly so a - * tenant with stale my.cnf overrides converges. - */ -final class MysqlTuningCatalog { - - static final String STATS_PERSISTENT = "STATS_PERSISTENT"; - static final String STATS_AUTO_RECALC = "STATS_AUTO_RECALC"; - static final String STATS_SAMPLE_PAGES = "STATS_SAMPLE_PAGES"; - - record Profile(Map settings, long rowThreshold, boolean relax, String reason) { - Profile { - settings = Map.copyOf(settings); - } - } - - private static final Map HOT = - Map.of( - STATS_PERSISTENT, "1", - STATS_AUTO_RECALC, "1", - STATS_SAMPLE_PAGES, "100"); - - private static final Map ENTITY_LARGE = - Map.of( - STATS_PERSISTENT, "1", - STATS_AUTO_RECALC, "1", - STATS_SAMPLE_PAGES, "64"); - - private static final Map ENTITY_SERVICE = - Map.of( - STATS_PERSISTENT, "1", - STATS_AUTO_RECALC, "1", - STATS_SAMPLE_PAGES, "32"); - - private static final long ROW_THRESHOLD_HOT = 0; - private static final long ROW_THRESHOLD_ENTITY_LARGE = 10_000; - private static final long ROW_THRESHOLD_ENTITY_SERVICE = 5_000; - - private static final Map CATALOG = buildCatalog(); - - private MysqlTuningCatalog() {} - - static Map catalog() { - return CATALOG; - } - - static Set tableNames() { - return CATALOG.keySet(); - } - - static Profile profileFor(final String tableName) { - return CATALOG.get(tableName); - } - - private static Map buildCatalog() { - Map map = new LinkedHashMap<>(); - map.put( - "entity_relationship", - new Profile(HOT, ROW_THRESHOLD_HOT, false, "Join target; raise sampling for planner")); - map.put("tag_usage", new Profile(HOT, ROW_THRESHOLD_HOT, false, "Hottest table on read path")); - addEntityLarge(map); - addEntityService(map); - return Map.copyOf(map); - } - - private static void addEntityLarge(final Map map) { - String reason = "Large entity table; bump InnoDB stats sampling"; - for (String t : - new String[] { - "storage_container_entity", - "table_entity", - "dashboard_entity", - "pipeline_entity", - "chart_entity", - "topic_entity", - "ml_model_entity", - "glossary_term_entity", - "metric_entity", - "report_entity", - "search_index_entity", - "api_collection_entity", - "api_endpoint_entity", - "dashboard_data_model_entity", - "ingestion_pipeline_entity", - "data_contract_entity", - "stored_procedure_entity", - "directory_entity", - "file_entity", - "spreadsheet_entity", - "worksheet_entity", - "query_entity" - }) { - map.put(t, new Profile(ENTITY_LARGE, ROW_THRESHOLD_ENTITY_LARGE, false, reason)); - } - } - - private static void addEntityService(final Map map) { - String reason = "Service-tier table; mild stats sampling bump"; - map.put( - "database_entity", - new Profile(ENTITY_SERVICE, ROW_THRESHOLD_ENTITY_SERVICE, false, reason)); - map.put( - "database_schema_entity", - new Profile(ENTITY_SERVICE, ROW_THRESHOLD_ENTITY_SERVICE, false, reason)); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresAutoTuner.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresAutoTuner.java deleted file mode 100644 index b1b002738209..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresAutoTuner.java +++ /dev/null @@ -1,281 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.ArrayList; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Locale; -import java.util.Map; -import java.util.stream.Collectors; -import org.jdbi.v3.core.Handle; -import org.openmetadata.service.util.dbtune.PostgresTuningCatalog.Profile; - -public final class PostgresAutoTuner implements AutoTuner { - - private static final List RELOPTION_KEYS = - List.of( - PostgresTuningCatalog.AUTOVACUUM_VACUUM_SCALE_FACTOR, - PostgresTuningCatalog.AUTOVACUUM_ANALYZE_SCALE_FACTOR, - PostgresTuningCatalog.AUTOVACUUM_VACUUM_COST_LIMIT, - PostgresTuningCatalog.AUTOVACUUM_VACUUM_COST_DELAY); - - @Override - public DbTuneResult analyze(final Handle handle) { - String version = readVersion(handle); - List serverParams = readServerParams(handle); - List stats = loadTableStats(handle); - List recs = stats.stream().map(this::recommend).toList(); - return new DbTuneResult("PostgreSQL", version, serverParams, recs); - } - - @Override - public TableRecommendation recommend(final TableStats stats) { - Profile profile = PostgresTuningCatalog.profileFor(stats.tableName()); - if (profile == null) { - return skip(stats, "Table is not in the dbtune catalog"); - } - if (stats.rowCount() < profile.rowThreshold()) { - return skip( - stats, - String.format( - Locale.ROOT, - "Row count %d below threshold %d", - stats.rowCount(), - profile.rowThreshold())); - } - return decideAction(stats, profile); - } - - private TableRecommendation decideAction(final TableStats stats, final Profile profile) { - Map recommended = profile.settings(); - Map current = stats.currentSettings(); - if (settingsMatch(current, recommended)) { - return new TableRecommendation( - stats.tableName(), - Action.OK, - stats.rowCount(), - stats.totalBytes(), - current, - recommended, - "Already matches recommended settings"); - } - Action action = chooseAction(current, profile); - return new TableRecommendation( - stats.tableName(), - action, - stats.rowCount(), - stats.totalBytes(), - current, - recommended, - profile.reason()); - } - - private Action chooseAction(final Map current, final Profile profile) { - if (current.isEmpty()) { - return profile.relax() ? Action.RELAX : Action.APPLY; - } - return profile.relax() ? Action.RELAX : Action.TIGHTEN; - } - - @Override - public void apply(final Handle handle, final TableRecommendation recommendation) { - if (!recommendation.action().isActionable()) { - return; - } - handle.execute(buildAlterStatement(recommendation)); - } - - @Override - public void analyzeOne(final Handle handle, final String tableName) { - handle.execute("ANALYZE " + quoteIdent(tableName)); - } - - @Override - public String buildAlterStatement(final TableRecommendation recommendation) { - String settings = - recommendation.recommendedSettings().entrySet().stream() - .sorted(Map.Entry.comparingByKey()) - .map(e -> e.getKey() + " = " + e.getValue()) - .collect(Collectors.joining(", ")); - return "ALTER TABLE " + quoteIdent(recommendation.tableName()) + " SET (" + settings + ")"; - } - - // ---- DB I/O ---- - - String readVersion(final Handle handle) { - return handle.createQuery("SHOW server_version").mapTo(String.class).findOne().orElse(""); - } - - List loadTableStats(final Handle handle) { - List result = new ArrayList<>(); - for (String tableName : PostgresTuningCatalog.tableNames()) { - TableStats stats = loadTableStats(handle, tableName); - if (stats != null) { - result.add(stats); - } - } - return result; - } - - TableStats loadTableStats(final Handle handle, final String tableName) { - return handle - .createQuery( - "SELECT c.reltuples::bigint AS rows, " - + " pg_relation_size(c.oid) AS heap_bytes, " - + " pg_indexes_size(c.oid) AS idx_bytes, " - + " COALESCE(c.reloptions, ARRAY[]::text[]) AS opts " - + "FROM pg_class c " - + "JOIN pg_namespace n ON n.oid = c.relnamespace " - + "WHERE c.relkind = 'r' " - + " AND n.nspname = ANY (current_schemas(false)) " - + " AND c.relname = :name") - .bind("name", tableName) - .map( - (rs, ctx) -> { - long rows = rs.getLong("rows"); - long heap = rs.getLong("heap_bytes"); - long idx = rs.getLong("idx_bytes"); - String[] opts = (String[]) rs.getArray("opts").getArray(); - return new TableStats(tableName, Math.max(rows, 0), heap, idx, parseReloptions(opts)); - }) - .findOne() - .orElse(null); - } - - @Override - public Map currentSettingsForTable(final Handle handle, final String tableName) { - return handle - .createQuery( - "SELECT COALESCE(c.reloptions, ARRAY[]::text[]) AS opts " - + "FROM pg_class c " - + "JOIN pg_namespace n ON n.oid = c.relnamespace " - + "WHERE c.relkind = 'r' " - + " AND n.nspname = ANY (current_schemas(false)) " - + " AND c.relname = :name") - .bind("name", tableName) - .map((rs, ctx) -> parseReloptions((String[]) rs.getArray("opts").getArray())) - .findOne() - .orElse(Map.of()); - } - - List readServerParams(final Handle handle) { - List checks = new ArrayList<>(); - Map recommendations = recommendedServerParams(); - for (Map.Entry e : recommendations.entrySet()) { - String name = e.getKey(); - String recommended = e.getValue(); - String current = - handle - .createQuery("SELECT setting FROM pg_settings WHERE name = :n") - .bind("n", name) - .mapTo(String.class) - .findOne() - .orElse(null); - checks.add(buildServerCheck(name, current, recommended)); - } - return checks; - } - - // ---- helpers ---- - - static Map parseReloptions(final String[] opts) { - if (opts == null || opts.length == 0) { - return Map.of(); - } - Map out = new LinkedHashMap<>(); - for (String opt : opts) { - int eq = opt.indexOf('='); - if (eq > 0) { - String key = opt.substring(0, eq).toLowerCase(Locale.ROOT); - String value = opt.substring(eq + 1); - if (RELOPTION_KEYS.contains(key)) { - out.put(key, value); - } - } - } - return Map.copyOf(out); - } - - static boolean settingsMatch(final Map current, final Map rec) { - for (Map.Entry e : rec.entrySet()) { - String currentValue = current.get(e.getKey()); - if (currentValue == null || !numericEquals(currentValue, e.getValue())) { - return false; - } - } - return true; - } - - private static boolean numericEquals(final String a, final String b) { - try { - return Double.parseDouble(a) == Double.parseDouble(b); - } catch (NumberFormatException ex) { - return a.equals(b); - } - } - - private static TableRecommendation skip(final TableStats stats, final String reason) { - return new TableRecommendation( - stats.tableName(), - Action.SKIP, - stats.rowCount(), - stats.totalBytes(), - stats.currentSettings(), - Map.of(), - reason); - } - - static String quoteIdent(final String identifier) { - if (!identifier.matches("[a-zA-Z_][a-zA-Z0-9_]*")) { - throw new IllegalArgumentException( - "Refusing to build SQL with unsafe identifier: " + identifier); - } - return "\"" + identifier + "\""; - } - - /** Server-level recommendations from the production runbook. */ - static Map recommendedServerParams() { - Map map = new LinkedHashMap<>(); - map.put("shared_buffers", "40% of RAM (use formula form on RDS)"); - map.put("effective_cache_size", "75% of RAM (use formula form on RDS)"); - map.put("work_mem", "131072"); // 128 MB - map.put("maintenance_work_mem", "2097152"); // 2 GB - map.put("random_page_cost", "1.1"); - map.put("effective_io_concurrency", "200"); - map.put("max_parallel_workers_per_gather", "4"); - map.put("autovacuum_naptime", "15"); - map.put("autovacuum_vacuum_scale_factor", "0.05"); - map.put("autovacuum_analyze_scale_factor", "0.02"); - return Map.copyOf(map); - } - - static ServerParamCheck buildServerCheck( - final String name, final String current, final String recommended) { - if (current == null) { - return new ServerParamCheck( - name, "", recommended, ServerParamCheck.STATUS_UNKNOWN, "Parameter not visible"); - } - if (recommended.contains("%")) { - return new ServerParamCheck( - name, - current, - recommended, - ServerParamCheck.STATUS_UNTUNED, - "RAM-relative; verify in RDS"); - } - boolean ok = numericEquals(current, recommended); - String status = ok ? ServerParamCheck.STATUS_OK : ServerParamCheck.STATUS_MISMATCH; - return new ServerParamCheck(name, current, recommended, status, ""); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresDiagnostic.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresDiagnostic.java deleted file mode 100644 index be989e2ec025..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresDiagnostic.java +++ /dev/null @@ -1,269 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.ArrayList; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import lombok.extern.slf4j.Slf4j; -import org.jdbi.v3.core.Handle; - -/** - * Postgres diagnostic. Each finding category is queried in its own try block so that a missing - * extension or a stat view permission issue surfaces as a {@link DbTuneDiagnosis#notes()} entry - * rather than aborting the whole run. - * - *

Thresholds are baked in for v1; if operators want them tunable later they become CLI flags. - */ -@Slf4j -public final class PostgresDiagnostic implements Diagnostic { - - static final long UNUSED_INDEX_SIZE_BYTES = 10L * 1024 * 1024; - static final double DEAD_TUPLE_RATIO = 0.2; - static final long DEAD_TUPLE_MIN_LIVE_ROWS = 10_000; - static final double LOW_CACHE_HIT_RATIO = 0.9; - static final long LOW_CACHE_HIT_MIN_READS = 1_000; - static final int STALE_STATS_DAYS = 14; - static final long STALE_STATS_MIN_LIVE_ROWS = 1_000; - static final long SEQ_SCAN_RATIO = 10; - static final long SEQ_SCAN_MIN = 1_000; - static final int SLOW_QUERY_LIMIT = 10; - static final long SLOW_QUERY_MIN_CALLS = 100; - static final int QUERY_TRUNCATE = 100; - - @Override - public DbTuneDiagnosis diagnose(final Handle handle) { - List findings = new ArrayList<>(); - List notes = new ArrayList<>(); - runCategory(handle, notes, "unused indexes", h -> findings.addAll(unusedIndexes(h))); - runCategory(handle, notes, "dead tuples", h -> findings.addAll(highDeadTuples(h))); - runCategory(handle, notes, "cache hit", h -> findings.addAll(lowCacheHit(h))); - runCategory(handle, notes, "stale stats", h -> findings.addAll(staleStats(h))); - runCategory(handle, notes, "seq scans", h -> findings.addAll(seqScanHeavy(h))); - runCategory(handle, notes, "slow queries", h -> findings.addAll(slowQueries(h, notes))); - return new DbTuneDiagnosis(findings, notes); - } - - private void runCategory( - final Handle handle, - final List notes, - final String label, - final java.util.function.Consumer body) { - try { - body.accept(handle); - } catch (Exception e) { - LOG.warn("Diagnostic [{}] failed: {}", label, e.getMessage()); - notes.add(label + ": " + e.getMessage()); - } - } - - // ---- categories ---- - - List unusedIndexes(final Handle handle) { - return handle - .createQuery( - "SELECT s.schemaname, s.relname AS table_name, s.indexrelname AS index_name, " - + " s.idx_scan AS scans, " - + " pg_relation_size(s.indexrelid) AS bytes " - + "FROM pg_stat_user_indexes s " - + "JOIN pg_index i ON i.indexrelid = s.indexrelid " - + "WHERE s.idx_scan = 0 " - + " AND NOT i.indisunique " - + " AND NOT i.indisprimary " - + " AND pg_relation_size(s.indexrelid) > :min_bytes " - + "ORDER BY pg_relation_size(s.indexrelid) DESC " - + "LIMIT 50") - .bind("min_bytes", UNUSED_INDEX_SIZE_BYTES) - .map( - (rs, ctx) -> - new Finding( - DiagnosticCategory.UNUSED_INDEX, - Severity.WARN, - Map.of( - "table", rs.getString("table_name"), - "index", rs.getString("index_name"), - "size", DbTuneReport.formatBytes(rs.getLong("bytes")), - "scans", String.valueOf(rs.getLong("scans"))))) - .list(); - } - - List highDeadTuples(final Handle handle) { - return handle - .createQuery( - "SELECT relname AS table_name, " - + " n_live_tup, " - + " n_dead_tup, " - + " ROUND((n_dead_tup::numeric / GREATEST(n_live_tup, 1)) * 100, 2) AS dead_pct, " - + " last_autovacuum " - + "FROM pg_stat_user_tables " - + "WHERE n_live_tup > :min_live " - + " AND n_dead_tup::numeric / GREATEST(n_live_tup, 1) > :threshold " - + "ORDER BY n_dead_tup DESC " - + "LIMIT 25") - .bind("min_live", DEAD_TUPLE_MIN_LIVE_ROWS) - .bind("threshold", DEAD_TUPLE_RATIO) - .map( - (rs, ctx) -> - new Finding( - DiagnosticCategory.HIGH_DEAD_TUPLES, - Severity.WARN, - Map.of( - "table", rs.getString("table_name"), - "live_rows", String.valueOf(rs.getLong("n_live_tup")), - "dead_rows", String.valueOf(rs.getLong("n_dead_tup")), - "dead_ratio", rs.getString("dead_pct") + "%", - "last_vacuum", nullSafe(rs.getString("last_autovacuum"))))) - .list(); - } - - List lowCacheHit(final Handle handle) { - return handle - .createQuery( - "SELECT relname AS table_name, " - + " heap_blks_read, " - + " heap_blks_hit, " - + " ROUND(heap_blks_hit::numeric / NULLIF(heap_blks_hit + heap_blks_read, 0) * 100, 2) AS hit_pct " - + "FROM pg_statio_user_tables " - + "WHERE heap_blks_read > :min_reads " - + " AND heap_blks_hit::numeric / NULLIF(heap_blks_hit + heap_blks_read, 0) < :threshold " - + "ORDER BY heap_blks_read DESC " - + "LIMIT 25") - .bind("min_reads", LOW_CACHE_HIT_MIN_READS) - .bind("threshold", LOW_CACHE_HIT_RATIO) - .map( - (rs, ctx) -> - new Finding( - DiagnosticCategory.LOW_CACHE_HIT, - Severity.INFO, - Map.of( - "table", rs.getString("table_name"), - "heap_reads", String.valueOf(rs.getLong("heap_blks_read")), - "heap_hits", String.valueOf(rs.getLong("heap_blks_hit")), - "hit_pct", rs.getString("hit_pct") + "%"))) - .list(); - } - - List staleStats(final Handle handle) { - return handle - .createQuery( - "SELECT relname AS table_name, " - + " n_live_tup, " - + " COALESCE(last_autoanalyze, last_analyze) AS last_analyzed " - + "FROM pg_stat_user_tables " - + "WHERE n_live_tup > :min_live " - + " AND (COALESCE(last_autoanalyze, last_analyze) IS NULL " - + " OR COALESCE(last_autoanalyze, last_analyze) < now() - (:days || ' days')::interval) " - + "ORDER BY n_live_tup DESC " - + "LIMIT 25") - .bind("min_live", STALE_STATS_MIN_LIVE_ROWS) - .bind("days", STALE_STATS_DAYS) - .map( - (rs, ctx) -> - new Finding( - DiagnosticCategory.STALE_STATS, - Severity.WARN, - Map.of( - "table", rs.getString("table_name"), - "live_rows", String.valueOf(rs.getLong("n_live_tup")), - "last_analyzed", nullSafe(rs.getString("last_analyzed"))))) - .list(); - } - - List seqScanHeavy(final Handle handle) { - // Includes idx_scan=0 tables — those are the *worst* candidates for a missing index, not - // edge cases to filter out. NULLIF would silently drop them via NULL comparison. - return handle - .createQuery( - "SELECT relname AS table_name, seq_scan, idx_scan " - + "FROM pg_stat_user_tables " - + "WHERE seq_scan > :min_seq " - + " AND (idx_scan = 0 OR seq_scan::numeric / idx_scan > :ratio) " - + "ORDER BY seq_scan DESC " - + "LIMIT 25") - .bind("min_seq", SEQ_SCAN_MIN) - .bind("ratio", SEQ_SCAN_RATIO) - .map( - (rs, ctx) -> - new Finding( - DiagnosticCategory.SEQ_SCAN_HEAVY, - Severity.INFO, - Map.of( - "table", rs.getString("table_name"), - "seq_scans", String.valueOf(rs.getLong("seq_scan")), - "idx_scans", String.valueOf(rs.getLong("idx_scan")), - "ratio", - formatSeqIdxRatio(rs.getLong("seq_scan"), rs.getLong("idx_scan"))))) - .list(); - } - - List slowQueries(final Handle handle, final List notes) { - if (!hasPgStatStatements(handle)) { - notes.add("slow queries: pg_stat_statements extension not installed"); - return List.of(); - } - return handle - .createQuery( - "SELECT query, calls, mean_exec_time AS mean_ms " - + "FROM pg_stat_statements " - + "WHERE calls > :min_calls " - + "ORDER BY mean_exec_time DESC " - + "LIMIT :limit") - .bind("min_calls", SLOW_QUERY_MIN_CALLS) - .bind("limit", SLOW_QUERY_LIMIT) - .map( - (rs, ctx) -> { - Map attrs = new LinkedHashMap<>(); - attrs.put("query", truncate(rs.getString("query"))); - attrs.put("calls", String.valueOf(rs.getLong("calls"))); - attrs.put( - "mean_ms", String.format(java.util.Locale.ROOT, "%.1f", rs.getDouble("mean_ms"))); - return new Finding(DiagnosticCategory.SLOW_QUERY, Severity.INFO, attrs); - }) - .list(); - } - - private boolean hasPgStatStatements(final Handle handle) { - return handle - .createQuery("SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements'") - .mapTo(Integer.class) - .findOne() - .isPresent(); - } - - static String truncate(final String query) { - if (query == null) { - return ""; - } - String collapsed = query.replaceAll("\\s+", " ").trim(); - return collapsed.length() <= QUERY_TRUNCATE - ? collapsed - : collapsed.substring(0, QUERY_TRUNCATE) + "…"; - } - - /** Empty string for SQL NULL — never the literal "null" since this lands in user-facing output. */ - static String nullSafe(final String value) { - return value == null ? "" : value; - } - - /** - * Formats {@code seq_scan / idx_scan} as a one-decimal ratio (e.g. {@code 7.5}) using {@code - * double} division. Returns {@code "∞"} when {@code idx_scan == 0}. - */ - static String formatSeqIdxRatio(final long seqScan, final long idxScan) { - if (idxScan == 0) { - return "∞"; - } - return String.format(java.util.Locale.ROOT, "%.1f", (double) seqScan / idxScan); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresTuningCatalog.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresTuningCatalog.java deleted file mode 100644 index df2b7e0e7705..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/PostgresTuningCatalog.java +++ /dev/null @@ -1,143 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.Set; - -/** - * Static catalog of which tables get which Postgres autovacuum reloptions, and the row-count - * threshold below which we skip tuning that table (small dev installs don't need aggressive - * autovacuum). Values come from production analysis of the 600k-container tenant. - */ -final class PostgresTuningCatalog { - - static final String AUTOVACUUM_VACUUM_SCALE_FACTOR = "autovacuum_vacuum_scale_factor"; - static final String AUTOVACUUM_ANALYZE_SCALE_FACTOR = "autovacuum_analyze_scale_factor"; - static final String AUTOVACUUM_VACUUM_COST_LIMIT = "autovacuum_vacuum_cost_limit"; - static final String AUTOVACUUM_VACUUM_COST_DELAY = "autovacuum_vacuum_cost_delay"; - - /** A tuning recipe for one table. */ - record Profile(Map settings, long rowThreshold, boolean relax, String reason) { - Profile { - settings = Map.copyOf(settings); - } - } - - private static final Map HOT_RELATIONSHIP = - Map.of( - AUTOVACUUM_ANALYZE_SCALE_FACTOR, "0.005", - AUTOVACUUM_VACUUM_SCALE_FACTOR, "0.01", - AUTOVACUUM_VACUUM_COST_LIMIT, "4000"); - - private static final Map HOT_TAG_USAGE = - Map.of( - AUTOVACUUM_ANALYZE_SCALE_FACTOR, "0.005", - AUTOVACUUM_VACUUM_SCALE_FACTOR, "0.01", - AUTOVACUUM_VACUUM_COST_LIMIT, "4000", - AUTOVACUUM_VACUUM_COST_DELAY, "0"); - - private static final Map ENTITY_LARGE = - Map.of( - AUTOVACUUM_ANALYZE_SCALE_FACTOR, "0.01", - AUTOVACUUM_VACUUM_SCALE_FACTOR, "0.02"); - - private static final Map ENTITY_SERVICE = - Map.of( - AUTOVACUUM_ANALYZE_SCALE_FACTOR, "0.02", - AUTOVACUUM_VACUUM_SCALE_FACTOR, "0.05"); - - private static final Map APPEND_ONLY = - Map.of( - AUTOVACUUM_ANALYZE_SCALE_FACTOR, "0.1", - AUTOVACUUM_VACUUM_SCALE_FACTOR, "0.2"); - - private static final long ROW_THRESHOLD_HOT = 0; - private static final long ROW_THRESHOLD_ENTITY_LARGE = 10_000; - private static final long ROW_THRESHOLD_ENTITY_SERVICE = 5_000; - private static final long ROW_THRESHOLD_APPEND_ONLY = 50_000; - - private static final Map CATALOG = buildCatalog(); - - private PostgresTuningCatalog() {} - - static Map catalog() { - return CATALOG; - } - - static Set tableNames() { - return CATALOG.keySet(); - } - - static Profile profileFor(final String tableName) { - return CATALOG.get(tableName); - } - - private static Map buildCatalog() { - Map map = new LinkedHashMap<>(); - map.put( - "entity_relationship", - new Profile(HOT_RELATIONSHIP, ROW_THRESHOLD_HOT, false, "Join target, write-heavy")); - map.put( - "tag_usage", - new Profile(HOT_TAG_USAGE, ROW_THRESHOLD_HOT, false, "Hottest table on read path")); - addEntityLarge(map); - addEntityService(map); - map.put( - "change_event", - new Profile(APPEND_ONLY, ROW_THRESHOLD_APPEND_ONLY, true, "Append-only, relax autovacuum")); - return Map.copyOf(map); - } - - private static void addEntityLarge(final Map map) { - String reason = "Large entity table; tighten autovacuum so list count stats stay fresh"; - for (String t : - new String[] { - "storage_container_entity", - "table_entity", - "dashboard_entity", - "pipeline_entity", - "chart_entity", - "topic_entity", - "ml_model_entity", - "glossary_term_entity", - "metric_entity", - "report_entity", - "search_index_entity", - "api_collection_entity", - "api_endpoint_entity", - "dashboard_data_model_entity", - "ingestion_pipeline_entity", - "data_contract_entity", - "stored_procedure_entity", - "directory_entity", - "file_entity", - "spreadsheet_entity", - "worksheet_entity", - "query_entity" - }) { - map.put(t, new Profile(ENTITY_LARGE, ROW_THRESHOLD_ENTITY_LARGE, false, reason)); - } - } - - private static void addEntityService(final Map map) { - String reason = "Service-tier table; mild tightening"; - map.put( - "database_entity", - new Profile(ENTITY_SERVICE, ROW_THRESHOLD_ENTITY_SERVICE, false, reason)); - map.put( - "database_schema_entity", - new Profile(ENTITY_SERVICE, ROW_THRESHOLD_ENTITY_SERVICE, false, reason)); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/ServerParamCheck.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/ServerParamCheck.java deleted file mode 100644 index 39142703e3bd..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/ServerParamCheck.java +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -public record ServerParamCheck( - String parameter, String currentValue, String recommendedValue, String status, String note) { - - public static final String STATUS_OK = "OK"; - - /** - * Direction-agnostic. Some recommended values (e.g. {@code random_page_cost = 1.1}, - * {@code autovacuum_*_scale_factor}) are deliberately lower than the engine default — labelling - * those mismatches as "undersized" would be wrong. Operators see the actual current vs - * recommended values in the report and can judge direction themselves. - */ - public static final String STATUS_MISMATCH = "MISMATCH"; - - public static final String STATUS_UNTUNED = "UNTUNED"; - public static final String STATUS_UNKNOWN = "UNKNOWN"; -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Severity.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Severity.java deleted file mode 100644 index 05993ef669d4..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/Severity.java +++ /dev/null @@ -1,18 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -public enum Severity { - INFO, - WARN -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/TableRecommendation.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/TableRecommendation.java deleted file mode 100644 index ecf509dac9be..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/TableRecommendation.java +++ /dev/null @@ -1,30 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.Map; - -public record TableRecommendation( - String tableName, - Action action, - long rowCount, - long totalBytes, - Map currentSettings, - Map recommendedSettings, - String reason) { - - public TableRecommendation { - currentSettings = currentSettings == null ? Map.of() : Map.copyOf(currentSettings); - recommendedSettings = recommendedSettings == null ? Map.of() : Map.copyOf(recommendedSettings); - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/TableStats.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/TableStats.java deleted file mode 100644 index 9a8d8ec4e421..000000000000 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/dbtune/TableStats.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import java.util.Map; - -public record TableStats( - String tableName, - long rowCount, - long dataBytes, - long indexBytes, - Map currentSettings) { - - public TableStats { - currentSettings = currentSettings == null ? Map.of() : Map.copyOf(currentSettings); - } - - public long totalBytes() { - return dataBytes + indexBytes; - } -} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/ReindexingUtil.java b/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/ReindexingUtil.java index 60a7edf0ecaf..34d49476a614 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/ReindexingUtil.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/workflows/searchIndex/ReindexingUtil.java @@ -13,7 +13,7 @@ package org.openmetadata.service.workflows.searchIndex; -import static org.openmetadata.service.apps.bundles.searchIndex.SearchIndexEntityTypes.TIME_SERIES_ENTITIES; +import static org.openmetadata.service.apps.bundles.searchIndex.SearchIndexApp.TIME_SERIES_ENTITIES; import static org.openmetadata.service.search.SearchClient.GLOBAL_SEARCH_ALIAS; import com.fasterxml.jackson.databind.JsonNode; diff --git a/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json b/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json index a44897895d3b..781b1a047c9c 100644 --- a/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json +++ b/openmetadata-service/src/main/resources/json/data/app/SearchIndexingApplication.json @@ -5,6 +5,7 @@ "entities": [ "all" ], + "recreateIndex": true, "batchSize": "100", "payLoadSize": 104857600, "producerThreads": 1, diff --git a/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/SearchIndexingApplication.json b/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/SearchIndexingApplication.json index 8fae312b7bd4..b45855c08546 100644 --- a/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/SearchIndexingApplication.json +++ b/openmetadata-service/src/main/resources/json/data/appMarketPlaceDefinition/SearchIndexingApplication.json @@ -2,7 +2,7 @@ "name": "SearchIndexingApplication", "displayName": "Search Indexing", "description": "OpenMetadata connects with Elastic/Open Search to provide search feature for Data Assets. This application provides additional features related to ES/OS.", - "features": "Sync OpenMetadata and Elastic Search with staged index promotion.", + "features": "Sync OpenMetadata and Elastic Search and Recreate Indexes.", "appType": "internal", "appScreenshots": ["SearchIndexPic1"], "developer": "Collate Inc.", @@ -20,6 +20,7 @@ "entities": [ "all" ], + "recreateIndex": false, "batchSize": "100", "payLoadSize": 104857600, "producerThreads": 1, diff --git a/openmetadata-service/src/main/resources/rdf/inference-rules/domain-membership-inheritance.json b/openmetadata-service/src/main/resources/rdf/inference-rules/domain-membership-inheritance.json new file mode 100644 index 000000000000..8f7c332513b6 --- /dev/null +++ b/openmetadata-service/src/main/resources/rdf/inference-rules/domain-membership-inheritance.json @@ -0,0 +1,10 @@ +{ + "name": "domain-membership-inheritance", + "displayName": "Domain membership inheritance", + "description": "If a Table belongs to a Domain via om:belongsToDomain, every Column of that Table inherits the same domain membership. Lets SPARQL queries against om:belongsToDomain return both table-level and column-level results without separate lookups.", + "ruleType": "CONSTRUCT", + "priority": 400, + "enabled": true, + "tags": ["governance"], + "ruleBody": "PREFIX om: \nCONSTRUCT { ?column om:belongsToDomain ?domain }\nWHERE {\n ?table a om:Table ;\n om:belongsToDomain ?domain ;\n om:hasColumn ?column .\n}" +} diff --git a/openmetadata-service/src/main/resources/rdf/inference-rules/pii-propagation-via-lineage.json b/openmetadata-service/src/main/resources/rdf/inference-rules/pii-propagation-via-lineage.json new file mode 100644 index 000000000000..cb7dbf2d4890 --- /dev/null +++ b/openmetadata-service/src/main/resources/rdf/inference-rules/pii-propagation-via-lineage.json @@ -0,0 +1,10 @@ +{ + "name": "pii-propagation-via-lineage", + "displayName": "PII propagation through lineage", + "description": "If a column is tagged with a PII classification AND another column receives data from it via column-level lineage, propagate the same PII tag to the downstream column. The propagated tag is marked om:inferred so admins can distinguish derived tags from manually applied ones.", + "ruleType": "CONSTRUCT", + "priority": 200, + "enabled": true, + "tags": ["security", "lineage"], + "ruleBody": "PREFIX om: \nCONSTRUCT {\n ?downstream om:hasTag ?piiTag .\n ?downstream om:inferredTagSource ?upstream\n}\nWHERE {\n ?upstream om:hasTag ?piiTag .\n ?piiTag om:tagFQN ?fqn .\n FILTER(STRSTARTS(?fqn, \"PII.\"))\n ?colLineage om:fromColumn ?upstream ;\n om:toColumn ?downstream .\n}" +} diff --git a/openmetadata-service/src/main/resources/rdf/inference-rules/schema-tag-inheritance.json b/openmetadata-service/src/main/resources/rdf/inference-rules/schema-tag-inheritance.json new file mode 100644 index 000000000000..b6c77af81c73 --- /dev/null +++ b/openmetadata-service/src/main/resources/rdf/inference-rules/schema-tag-inheritance.json @@ -0,0 +1,10 @@ +{ + "name": "schema-tag-inheritance", + "displayName": "Schema → table → column tag inheritance", + "description": "Propagate tags down the containment hierarchy: a tag on a DatabaseSchema is inherited by every Table in that schema, and a tag on a Table is inherited by every Column. Inferred tags are marked om:inferredTagSource pointing to the parent that supplied them.", + "ruleType": "CONSTRUCT", + "priority": 300, + "enabled": true, + "tags": ["governance"], + "ruleBody": "PREFIX om: \nCONSTRUCT {\n ?descendant om:hasTag ?tag .\n ?descendant om:inferredTagSource ?ancestor\n}\nWHERE {\n {\n ?ancestor a om:DatabaseSchema ;\n om:hasTag ?tag .\n ?descendant om:belongsToSchema ?ancestor .\n } UNION {\n ?ancestor a om:Table ;\n om:hasTag ?tag .\n ?ancestor om:hasColumn ?descendant .\n }\n}" +} diff --git a/openmetadata-service/src/main/resources/rdf/inference-rules/transitive-lineage-closure.json b/openmetadata-service/src/main/resources/rdf/inference-rules/transitive-lineage-closure.json new file mode 100644 index 000000000000..7bd0e6ee5805 --- /dev/null +++ b/openmetadata-service/src/main/resources/rdf/inference-rules/transitive-lineage-closure.json @@ -0,0 +1,10 @@ +{ + "name": "transitive-lineage-closure", + "displayName": "Transitive lineage closure", + "description": "Materialize indirect upstream/downstream lineage edges by walking prov:wasDerivedFrom transitively. Lets SPARQL answer 'all upstream tables of dashboard X' without requiring property paths in user queries.", + "ruleType": "CONSTRUCT", + "priority": 100, + "enabled": true, + "tags": ["lineage"], + "ruleBody": "PREFIX prov: \nPREFIX om: \nCONSTRUCT { ?x om:transitivelyDerivedFrom ?y }\nWHERE {\n ?x prov:wasDerivedFrom+ ?y .\n FILTER(?x != ?y)\n}" +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoffTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoffTest.java new file mode 100644 index 000000000000..5906b152bf25 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/AdaptiveBackoffTest.java @@ -0,0 +1,72 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +@DisplayName("AdaptiveBackoff Tests") +class AdaptiveBackoffTest { + + @Test + @DisplayName("returns initial delay on first call") + void initialDelay() { + AdaptiveBackoff backoff = new AdaptiveBackoff(100, 2000); + assertEquals(100, backoff.nextDelay()); + } + + @Test + @DisplayName("doubles delay on each subsequent call") + void exponentialDoubling() { + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 10000); + assertEquals(50, backoff.nextDelay()); + assertEquals(100, backoff.nextDelay()); + assertEquals(200, backoff.nextDelay()); + assertEquals(400, backoff.nextDelay()); + assertEquals(800, backoff.nextDelay()); + } + + @Test + @DisplayName("caps at maxMs") + void capAtMax() { + AdaptiveBackoff backoff = new AdaptiveBackoff(100, 300); + assertEquals(100, backoff.nextDelay()); + assertEquals(200, backoff.nextDelay()); + assertEquals(300, backoff.nextDelay()); + assertEquals(300, backoff.nextDelay()); + } + + @Test + @DisplayName("reset returns to initial delay") + void resetToInitial() { + AdaptiveBackoff backoff = new AdaptiveBackoff(50, 1000); + backoff.nextDelay(); + backoff.nextDelay(); + backoff.nextDelay(); + + backoff.reset(); + assertEquals(50, backoff.nextDelay()); + } + + @Test + @DisplayName("rejects invalid initialMs") + void rejectsInvalidInitialMs() { + assertThrows(IllegalArgumentException.class, () -> new AdaptiveBackoff(0, 1000)); + assertThrows(IllegalArgumentException.class, () -> new AdaptiveBackoff(-1, 1000)); + } + + @Test + @DisplayName("rejects maxMs less than initialMs") + void rejectsMaxLessThanInitial() { + assertThrows(IllegalArgumentException.class, () -> new AdaptiveBackoff(200, 100)); + } + + @Test + @DisplayName("works when initialMs equals maxMs") + void initialEqualsMax() { + AdaptiveBackoff backoff = new AdaptiveBackoff(500, 500); + assertEquals(500, backoff.nextDelay()); + assertEquals(500, backoff.nextDelay()); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListenerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListenerTest.java index 2c1de34b3b75..6c278f6e8610 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListenerTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/CompositeProgressListenerTest.java @@ -97,6 +97,11 @@ public UUID getAppId() { return UUID.fromString("00000000-0000-0000-0000-000000000002"); } + @Override + public boolean isDistributed() { + return false; + } + @Override public String getSource() { return "UNIT_TEST"; diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategyTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategyTest.java index 042f19abf42e..804225e9cbda 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategyTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedIndexingStrategyTest.java @@ -23,7 +23,6 @@ import java.util.Map; import java.util.Set; import java.util.UUID; -import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -301,54 +300,6 @@ void updateStatsFromDistributedJobFallsBackToLocalSinkAndPartitionStats() throws } } - @Test - void updateStatsFromDistributedJobUsesAggregatedServerStatsWhenOnlySinkFailures() - throws Exception { - CollectionDAO.SearchIndexServerStatsDAO serverStatsDao = - mock(CollectionDAO.SearchIndexServerStatsDAO.class); - UUID jobId = UUID.fromString("00000000-0000-0000-0000-000000000023"); - Stats stats = createBaseStats("table", 10); - SearchIndexJob distributedJob = - SearchIndexJob.builder() - .id(jobId) - .totalRecords(10) - .successRecords(8) - .failedRecords(2) - .entityStats( - Map.of( - "table", - SearchIndexJob.EntityTypeStats.builder() - .entityType("table") - .totalRecords(10) - .successRecords(0) - .failedRecords(10) - .build())) - .build(); - - when(collectionDAO.searchIndexServerStatsDAO()).thenReturn(serverStatsDao); - when(serverStatsDao.getAggregatedStats(jobId.toString())) - .thenReturn( - new CollectionDAO.SearchIndexServerStatsDAO.AggregatedServerStats( - 10, 0, 0, 0, 10, 10, 0, 0, 0, 0, 0, 0, 0, 0, 1)); - - try (MockedStatic entityMock = mockStatic(Entity.class)) { - entityMock.when(Entity::getCollectionDAO).thenReturn(collectionDAO); - - invokePrivate( - "updateStatsFromDistributedJob", - new Class[] {Stats.class, SearchIndexJob.class, StepStats.class}, - stats, - distributedJob, - new StepStats().withSuccessRecords(8).withFailedRecords(2)); - } - - assertEquals(0, stats.getJobStats().getSuccessRecords()); - assertEquals(10, stats.getJobStats().getFailedRecords()); - assertEquals(10, stats.getSinkStats().getTotalRecords()); - assertEquals(0, stats.getSinkStats().getSuccessRecords()); - assertEquals(10, stats.getSinkStats().getFailedRecords()); - } - @Test void statusHelpersReportStoppedIncompleteAndCompleteJobs() throws Exception { Stats complete = createBaseStats("table", 10); @@ -381,47 +332,14 @@ void statusHelpersReportStoppedIncompleteAndCompleteJobs() throws Exception { } @Test - @SuppressWarnings("unchecked") - void finalizeAllEntityReindexPromotesZeroRecordEntityFromInitializedStats() throws Exception { - DistributedSearchIndexExecutor executor = mock(DistributedSearchIndexExecutor.class); - EntityCompletionTracker tracker = mock(EntityCompletionTracker.class); - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = stagedContext("user"); - - when(tracker.getPromotedEntities()).thenReturn(Set.of()); - when(executor.getEntityTracker()).thenReturn(tracker); - when(executor.getJobWithFreshStats()) - .thenReturn(SearchIndexJob.builder().entityStats(Map.of()).build()); - setField("distributedExecutor", executor); - ((AtomicReference) getField("currentStats")).set(createBaseStats("user", 0)); - - boolean result = - (Boolean) - invokePrivate( - "finalizeAllEntityReindex", - new Class[] {RecreateIndexHandler.class, ReindexContext.class, boolean.class}, - indexPromotionHandler, - stagedIndexContext, - true); - - assertTrue(result); - ArgumentCaptor contextCaptor = - ArgumentCaptor.forClass(EntityReindexContext.class); - ArgumentCaptor successCaptor = ArgumentCaptor.forClass(Boolean.class); - verify(indexPromotionHandler).finalizeReindex(contextCaptor.capture(), successCaptor.capture()); - assertEquals("user", contextCaptor.getValue().getEntityType()); - assertEquals(Boolean.TRUE, successCaptor.getValue()); - } - - @Test - void finalizeAllEntityReindexSkipsPromotedEntitiesAndFailsMissingEntityStats() throws Exception { + void finalizeAllEntityReindexSkipsPromotedEntitiesAndUsesPerEntitySuccess() throws Exception { DistributedSearchIndexExecutor executor = mock(DistributedSearchIndexExecutor.class); EntityCompletionTracker tracker = mock(EntityCompletionTracker.class); - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = new ReindexContext(); - stagedIndexContext.add( + RecreateIndexHandler recreateIndexHandler = mock(RecreateIndexHandler.class); + ReindexContext recreateContext = new ReindexContext(); + recreateContext.add( "table", "table_index", "table_original", "table_staged", Set.of(), "table", List.of()); - stagedIndexContext.add( + recreateContext.add( "user", "user_index", "user_original", @@ -429,7 +347,7 @@ void finalizeAllEntityReindexSkipsPromotedEntitiesAndFailsMissingEntityStats() t Set.of("user"), "user", List.of("parent")); - stagedIndexContext.add( + recreateContext.add( "dashboard", "dash_index", "dash_original", @@ -460,8 +378,8 @@ void finalizeAllEntityReindexSkipsPromotedEntitiesAndFailsMissingEntityStats() t invokePrivate( "finalizeAllEntityReindex", new Class[] {RecreateIndexHandler.class, ReindexContext.class, boolean.class}, - indexPromotionHandler, - stagedIndexContext, + recreateIndexHandler, + recreateContext, true); assertTrue(result); @@ -469,7 +387,7 @@ void finalizeAllEntityReindexSkipsPromotedEntitiesAndFailsMissingEntityStats() t ArgumentCaptor contextCaptor = ArgumentCaptor.forClass(EntityReindexContext.class); ArgumentCaptor successCaptor = ArgumentCaptor.forClass(Boolean.class); - verify(indexPromotionHandler, times(2)) + verify(recreateIndexHandler, times(2)) .finalizeReindex(contextCaptor.capture(), successCaptor.capture()); Map outcomes = new java.util.HashMap<>(); @@ -478,7 +396,7 @@ void finalizeAllEntityReindexSkipsPromotedEntitiesAndFailsMissingEntityStats() t contextCaptor.getAllValues().get(i).getEntityType(), successCaptor.getAllValues().get(i)); } - assertEquals(Boolean.FALSE, outcomes.get("user")); + assertEquals(Boolean.TRUE, outcomes.get("user")); assertEquals(Boolean.FALSE, outcomes.get("dashboard")); } @@ -532,22 +450,12 @@ void executeReturnsCompletedResultForSuccessfulSinglePassDistributedRun() { .failedRecords(0) .build())) .build(); - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = stagedContext(Entity.TABLE); - ReindexingConfiguration reindexConfig = - ReindexingConfiguration.builder() - .entities(Set.of(Entity.TABLE)) - .batchSize(25) - .maxConcurrentRequests(3) - .payloadSize(1024L) - .build(); + RecreateIndexHandler recreateIndexHandler = mock(RecreateIndexHandler.class); when(entityRepository.getDao()).thenReturn(entityDao); when(entityDao.listCount(any(ListFilter.class))).thenReturn(5); when(searchRepository.createBulkSink(anyInt(), anyInt(), anyLong())).thenReturn(bulkSink); - when(searchRepository.createReindexHandler()).thenReturn(indexPromotionHandler); - when(indexPromotionHandler.reCreateIndexes(reindexConfig.entities())) - .thenReturn(stagedIndexContext); + when(searchRepository.createReindexHandler()).thenReturn(recreateIndexHandler); when(bulkSink.getPendingVectorTaskCount()).thenReturn(0); when(bulkSink.flushAndAwait(60)).thenReturn(true); when(bulkSink.getStats()) @@ -569,7 +477,15 @@ void executeReturnsCompletedResultForSuccessfulSinglePassDistributedRun() { entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); entityMock.when(Entity::getCollectionDAO).thenReturn(collectionDAO); - ExecutionResult result = strategy.execute(reindexConfig, context(jobId)); + ExecutionResult result = + strategy.execute( + ReindexingConfiguration.builder() + .entities(Set.of(Entity.TABLE)) + .batchSize(25) + .maxConcurrentRequests(3) + .payloadSize(1024L) + .build(), + context(jobId)); assertEquals(ExecutionResult.Status.COMPLETED, result.status()); assertEquals(5, result.totalRecords()); @@ -581,81 +497,17 @@ void executeReturnsCompletedResultForSuccessfulSinglePassDistributedRun() { DistributedSearchIndexExecutor constructed = executorConstruction.constructed().getFirst(); verify(constructed).performStartupRecovery(); verify(constructed).setAppContext(APP_ID, 1234L); - verify(constructed).execute(bulkSink, stagedIndexContext, reindexConfig); - } - } - - @Test - @SuppressWarnings({"rawtypes", "unchecked"}) - void executeNormalizesLegacyEntityAliasesBeforeDistributedSetup() { - @SuppressWarnings("unchecked") - EntityTimeSeriesRepository timeSeriesRepository = mock(EntityTimeSeriesRepository.class); - EntityTimeSeriesDAO timeSeriesDao = mock(EntityTimeSeriesDAO.class); - BulkSink bulkSink = mock(BulkSink.class); - UUID jobId = UUID.fromString("00000000-0000-0000-0000-000000000032"); - SearchIndexJob completedJob = - SearchIndexJob.builder() - .id(jobId) - .status(IndexJobStatus.COMPLETED) - .totalRecords(5) - .successRecords(5) - .failedRecords(0) - .entityStats( - Map.of( - Entity.QUERY_COST_RECORD, - SearchIndexJob.EntityTypeStats.builder() - .entityType(Entity.QUERY_COST_RECORD) - .totalRecords(5) - .successRecords(5) - .failedRecords(0) - .build())) - .build(); - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = stagedContext(Entity.QUERY_COST_RECORD); - ReindexingConfiguration reindexConfig = - ReindexingConfiguration.builder() - .entities(Set.of(SearchIndexEntityTypes.QUERY_COST_RESULT)) - .build(); - - when(timeSeriesRepository.getTimeSeriesDao()).thenReturn(timeSeriesDao); - when(timeSeriesDao.listCount(any(ListFilter.class))).thenReturn(5); - when(searchRepository.createBulkSink(anyInt(), anyInt(), anyLong())).thenReturn(bulkSink); - when(searchRepository.createReindexHandler()).thenReturn(indexPromotionHandler); - when(indexPromotionHandler.reCreateIndexes(Set.of(Entity.QUERY_COST_RECORD))) - .thenReturn(stagedIndexContext); - when(bulkSink.getPendingVectorTaskCount()).thenReturn(0); - when(bulkSink.flushAndAwait(60)).thenReturn(true); - when(bulkSink.getStats()) - .thenReturn(new StepStats().withSuccessRecords(5).withFailedRecords(0)); - when(bulkSink.getVectorStats()).thenReturn(new StepStats().withTotalRecords(0)); - - try (MockedStatic entityMock = mockStatic(Entity.class); - MockedConstruction executorConstruction = - mockConstruction( - DistributedSearchIndexExecutor.class, - (mock, context) -> { - when(mock.createJob( - any(Set.class), any(EventPublisherJob.class), eq("admin"), any())) - .thenReturn(completedJob); - when(mock.getJobWithFreshStats()).thenReturn(completedJob); - })) { - entityMock - .when(() -> Entity.getEntityTimeSeriesRepository(Entity.QUERY_COST_RECORD)) - .thenReturn(timeSeriesRepository); - - ExecutionResult result = strategy.execute(reindexConfig, context(jobId)); - - assertEquals(ExecutionResult.Status.COMPLETED, result.status()); - DistributedSearchIndexExecutor constructed = executorConstruction.constructed().getFirst(); - ArgumentCaptor entityTypesCaptor = ArgumentCaptor.forClass(Set.class); verify(constructed) - .createJob( - entityTypesCaptor.capture(), - any(EventPublisherJob.class), - eq("admin"), - eq(reindexConfig)); - assertEquals(Set.of(Entity.QUERY_COST_RECORD), entityTypesCaptor.getValue()); - verify(indexPromotionHandler).reCreateIndexes(Set.of(Entity.QUERY_COST_RECORD)); + .execute( + bulkSink, + null, + false, + ReindexingConfiguration.builder() + .entities(Set.of(Entity.TABLE)) + .batchSize(25) + .maxConcurrentRequests(3) + .payloadSize(1024L) + .build()); } } @@ -697,15 +549,10 @@ void executeClosesSinkAndReturnsFailedWhenDoExecuteThrowsAndSinkCloseAlsoFails() EntityRepository entityRepository = mock(EntityRepository.class); EntityDAO entityDao = mock(EntityDAO.class); BulkSink bulkSink = mock(BulkSink.class); - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = stagedContext(Entity.TABLE); when(entityRepository.getDao()).thenReturn(entityDao); when(entityDao.listCount(any(ListFilter.class))).thenReturn(5); when(searchRepository.createBulkSink(anyInt(), anyInt(), anyLong())).thenReturn(bulkSink); - when(searchRepository.createReindexHandler()).thenReturn(indexPromotionHandler); - when(indexPromotionHandler.reCreateIndexes(Set.of(Entity.TABLE))) - .thenReturn(stagedIndexContext); doThrow(new RuntimeException("close failed")).when(bulkSink).close(); try (MockedStatic entityMock = mockStatic(Entity.class); @@ -719,10 +566,7 @@ void executeClosesSinkAndReturnsFailedWhenDoExecuteThrowsAndSinkCloseAlsoFails() SearchIndexJob.builder().id(UUID.randomUUID()).totalRecords(5).build()); org.mockito.Mockito.doThrow(new RuntimeException("execute failed")) .when(mock) - .execute( - any(BulkSink.class), - any(ReindexContext.class), - any(ReindexingConfiguration.class)); + .execute(any(), any(), eq(false), any()); })) { entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); @@ -747,15 +591,10 @@ void executeClosesSinkSuccessfullyWhenDoExecuteThrows() throws Exception { EntityRepository entityRepository = mock(EntityRepository.class); EntityDAO entityDao = mock(EntityDAO.class); BulkSink bulkSink = mock(BulkSink.class); - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = stagedContext(Entity.TABLE); when(entityRepository.getDao()).thenReturn(entityDao); when(entityDao.listCount(any(ListFilter.class))).thenReturn(5); when(searchRepository.createBulkSink(anyInt(), anyInt(), anyLong())).thenReturn(bulkSink); - when(searchRepository.createReindexHandler()).thenReturn(indexPromotionHandler); - when(indexPromotionHandler.reCreateIndexes(Set.of(Entity.TABLE))) - .thenReturn(stagedIndexContext); try (MockedStatic entityMock = mockStatic(Entity.class); MockedConstruction executorConstruction = @@ -768,10 +607,7 @@ void executeClosesSinkSuccessfullyWhenDoExecuteThrows() throws Exception { SearchIndexJob.builder().id(UUID.randomUUID()).totalRecords(5).build()); org.mockito.Mockito.doThrow(new RuntimeException("execute failed")) .when(mock) - .execute( - any(BulkSink.class), - any(ReindexContext.class), - any(ReindexingConfiguration.class)); + .execute(any(), any(), eq(false), any()); })) { entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); @@ -804,19 +640,6 @@ private Stats createBaseStats(String entityType, int totalRecords) { return stats; } - private ReindexContext stagedContext(String entityType) { - ReindexContext context = new ReindexContext(); - context.add( - entityType, - entityType + "_index", - entityType + "_original", - entityType + "_staged", - Set.of(), - entityType, - List.of()); - return context; - } - private Object invokePrivate(String methodName, Class[] parameterTypes, Object... args) throws Exception { Method method = DistributedIndexingStrategy.class.getDeclaredMethod(methodName, parameterTypes); @@ -858,6 +681,11 @@ public UUID getAppId() { return APP_ID; } + @Override + public boolean isDistributed() { + return true; + } + @Override public String getSource() { return "UNIT_TEST"; diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexFinalizerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexFinalizerTest.java deleted file mode 100644 index 62100526b05d..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/DistributedReindexFinalizerTest.java +++ /dev/null @@ -1,96 +0,0 @@ -package org.openmetadata.service.apps.bundles.searchIndex; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; - -import java.util.List; -import java.util.Map; -import java.util.Set; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; -import org.openmetadata.service.Entity; -import org.openmetadata.service.apps.bundles.searchIndex.distributed.SearchIndexJob; -import org.openmetadata.service.search.EntityReindexContext; -import org.openmetadata.service.search.RecreateIndexHandler; -import org.openmetadata.service.search.ReindexContext; - -class DistributedReindexFinalizerTest { - - @Test - void finalizeRemainingEntitiesPromotesColumnOnceWhenTableAndColumnRemain() { - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = stagedContext(Entity.TABLE, Entity.TABLE_COLUMN); - - DistributedReindexFinalizer finalizer = - new DistributedReindexFinalizer(indexPromotionHandler, stagedIndexContext); - finalizer.finalizeRemainingEntities(Set.of(), Map.of(Entity.TABLE, successfulStats()), true); - - ArgumentCaptor contextCaptor = - ArgumentCaptor.forClass(EntityReindexContext.class); - ArgumentCaptor successCaptor = ArgumentCaptor.forClass(Boolean.class); - verify(indexPromotionHandler, times(2)) - .finalizeReindex(contextCaptor.capture(), successCaptor.capture()); - - Map finalizations = finalizations(contextCaptor, successCaptor); - assertEquals(Set.of(Entity.TABLE, Entity.TABLE_COLUMN), finalizations.keySet()); - assertEquals(Boolean.TRUE, finalizations.get(Entity.TABLE)); - assertEquals(Boolean.TRUE, finalizations.get(Entity.TABLE_COLUMN)); - } - - @Test - void finalizeRemainingEntitiesDoesNotRepromoteAlreadyPromotedColumnWhenTableRemains() { - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); - ReindexContext stagedIndexContext = stagedContext(Entity.TABLE, Entity.TABLE_COLUMN); - - DistributedReindexFinalizer finalizer = - new DistributedReindexFinalizer(indexPromotionHandler, stagedIndexContext); - finalizer.finalizeRemainingEntities( - Set.of(Entity.TABLE_COLUMN), Map.of(Entity.TABLE, successfulStats()), true); - - ArgumentCaptor contextCaptor = - ArgumentCaptor.forClass(EntityReindexContext.class); - ArgumentCaptor successCaptor = ArgumentCaptor.forClass(Boolean.class); - verify(indexPromotionHandler, times(1)) - .finalizeReindex(contextCaptor.capture(), successCaptor.capture()); - - assertEquals(Entity.TABLE, contextCaptor.getValue().getEntityType()); - assertEquals(Boolean.TRUE, successCaptor.getValue()); - } - - private Map finalizations( - ArgumentCaptor contextCaptor, ArgumentCaptor successCaptor) { - List contexts = contextCaptor.getAllValues(); - List outcomes = successCaptor.getAllValues(); - return Map.of( - contexts.get(0).getEntityType(), - outcomes.get(0), - contexts.get(1).getEntityType(), - outcomes.get(1)); - } - - private SearchIndexJob.EntityTypeStats successfulStats() { - return SearchIndexJob.EntityTypeStats.builder() - .entityType(Entity.TABLE) - .totalRecords(1) - .successRecords(1) - .failedRecords(0) - .build(); - } - - private ReindexContext stagedContext(String... entities) { - ReindexContext context = new ReindexContext(); - for (String entity : entities) { - context.add( - entity, - entity + "_index", - entity + "_original", - entity + "_staged", - Set.of(entity), - entity, - List.of()); - } - return context; - } -} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimatorTest.java new file mode 100644 index 000000000000..5f54e4f9f63f --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityBatchSizeEstimatorTest.java @@ -0,0 +1,67 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +@DisplayName("EntityBatchSizeEstimator Tests") +class EntityBatchSizeEstimatorTest { + + @Test + @DisplayName("LARGE entities get smaller batch size") + void largeEntitiesGetSmallerBatch() { + int base = 200; + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("table", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("topic", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("dashboard", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("mlmodel", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("container", base)); + assertEquals(100, EntityBatchSizeEstimator.estimateBatchSize("storedProcedure", base)); + } + + @Test + @DisplayName("LARGE entities respect minimum batch size of 25") + void largeEntitiesRespectMinimum() { + assertEquals(25, EntityBatchSizeEstimator.estimateBatchSize("table", 40)); + assertEquals(25, EntityBatchSizeEstimator.estimateBatchSize("table", 10)); + } + + @Test + @DisplayName("SMALL entities get larger batch size") + void smallEntitiesGetLargerBatch() { + int base = 200; + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("user", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("team", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("bot", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("role", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("policy", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("tag", base)); + assertEquals(400, EntityBatchSizeEstimator.estimateBatchSize("classification", base)); + } + + @Test + @DisplayName("SMALL entities respect maximum batch size of 1000") + void smallEntitiesRespectMaximum() { + assertEquals(1000, EntityBatchSizeEstimator.estimateBatchSize("user", 600)); + assertEquals(1000, EntityBatchSizeEstimator.estimateBatchSize("user", 800)); + } + + @Test + @DisplayName("MEDIUM (unknown) entities get base batch size unchanged") + void mediumEntitiesUnchanged() { + int base = 200; + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("pipeline", base)); + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("database", base)); + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("glossaryTerm", base)); + assertEquals(base, EntityBatchSizeEstimator.estimateBatchSize("unknownEntity", base)); + } + + @Test + @DisplayName("handles zero and negative base batch size gracefully") + void handlesZeroAndNegative() { + assertEquals(0, EntityBatchSizeEstimator.estimateBatchSize("table", 0)); + assertTrue(EntityBatchSizeEstimator.estimateBatchSize("table", -1) < 0); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderLifecycleTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderLifecycleTest.java new file mode 100644 index 000000000000..eb9699104b69 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderLifecycleTest.java @@ -0,0 +1,241 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.Phaser; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedConstruction; +import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.type.Paging; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.util.RestUtil; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntitiesSource; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntityTimeSeriesSource; +import org.openmetadata.service.workflows.searchIndex.ReindexingUtil; + +class EntityReaderLifecycleTest { + + private ExecutorService producerExecutor; + private AtomicBoolean stopped; + private EntityReader reader; + + @BeforeEach + void setUp() { + producerExecutor = mock(ExecutorService.class); + stopped = new AtomicBoolean(false); + reader = new EntityReader(producerExecutor, stopped, 1, 0); + when(producerExecutor.submit(any(Runnable.class))) + .thenAnswer( + invocation -> { + ((Runnable) invocation.getArgument(0)).run(); + return mock(Future.class); + }); + } + + @Test + void readEntityReturnsZeroWhenNoRecordsExist() { + Phaser phaser = new Phaser(1); + + int submitted = + reader.readEntity( + "table", 0, 50, phaser, (entityType, batch, offset) -> fail("callback should not run")); + + assertEquals(0, submitted); + assertEquals(1, phaser.getRegisteredParties()); + verifyNoInteractions(producerExecutor); + } + + @Test + void readEntityProcessesSingleRegularEntityReaderUntilCursorExhausted() throws Exception { + Phaser phaser = new Phaser(1); + List offsets = new ArrayList<>(); + ResultList batch = mockResult(List.of("table-1", "table-2"), null, 0); + + try (MockedConstruction construction = + mockConstruction( + PaginatedEntitiesSource.class, + (mock, context) -> + when(mock.readNextKeyset(isNull())).thenReturn((ResultList) batch))) { + + int submitted = + reader.readEntity( + "table", 2, 10, phaser, (entityType, result, offset) -> offsets.add(offset)); + + assertEquals(1, submitted); + assertEquals(List.of(0), offsets); + assertEquals(1, phaser.getRegisteredParties()); + assertEquals(2, construction.constructed().size()); + verify(construction.constructed().get(1)).readNextKeyset(null); + } + } + + @Test + void readEntityUsesTimeSeriesConstructorsAndBoundaryCursorsForParallelReaders() throws Exception { + Phaser phaser = new Phaser(1); + String entityType = ReportData.ReportDataType.ENTITY_REPORT_DATA.value(); + AtomicInteger callbackCount = new AtomicInteger(); + List> constructorArguments = new ArrayList<>(); + + try (MockedConstruction construction = + mockConstruction( + PaginatedEntityTimeSeriesSource.class, + (mock, context) -> { + constructorArguments.add(List.copyOf(context.arguments())); + when(mock.readWithCursor(any())) + .thenReturn((ResultList) mockResult(List.of("row"), null, 0)); + })) { + + int submitted = + reader.readEntity( + entityType, + 6, + 2, + phaser, + (type, result, offset) -> callbackCount.incrementAndGet(), + 100L, + 200L); + + assertEquals(3, submitted); + assertEquals(3, callbackCount.get()); + assertEquals(3, construction.constructed().size()); + assertEquals(1, phaser.getRegisteredParties()); + + assertEquals(List.of(entityType, 2, List.of(), 6, 100L, 200L), constructorArguments.get(0)); + assertEquals(List.of(entityType, 2, List.of(), 6, 100L, 200L), constructorArguments.get(1)); + assertEquals(List.of(entityType, 2, List.of(), 6, 100L, 200L), constructorArguments.get(2)); + + verify(construction.constructed().get(0)).readWithCursor(null); + verify(construction.constructed().get(1)).readWithCursor(RestUtil.encodeCursor("2")); + verify(construction.constructed().get(2)).readWithCursor(RestUtil.encodeCursor("4")); + } + } + + @Test + void readEntityDeregistersMissingReadersWhenBoundaryDiscoveryReturnsFewerCursors() { + Phaser phaser = new Phaser(1); + AtomicInteger constructionCount = new AtomicInteger(); + + try (MockedConstruction construction = + mockConstruction( + PaginatedEntitiesSource.class, + (mock, context) -> { + if (constructionCount.getAndIncrement() == 0) { + when(mock.findBoundaryCursors(anyInt(), anyInt())).thenReturn(List.of()); + } else { + when(mock.readNextKeyset(any())) + .thenReturn((ResultList) mockResult(List.of(), null, 0)); + } + })) { + + int submitted = + reader.readEntity( + "table", + 6, + 2, + phaser, + (entityType, batch, offset) -> fail("empty batch should not invoke callback")); + + assertEquals(3, submitted); + assertEquals(2, construction.constructed().size()); + assertEquals(1, phaser.getRegisteredParties()); + verify(producerExecutor).submit(any(Runnable.class)); + } + } + + @Test + void readEntityRestoresPhaserStateWhenSubmissionFails() { + Phaser phaser = new Phaser(1); + when(producerExecutor.submit(any(Runnable.class))) + .thenThrow(new IllegalStateException("submit failed")); + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> + reader.readEntity( + "table", + 2, + 10, + phaser, + (entityType, batch, offset) -> fail("callback should not run"))); + + assertEquals("submit failed", exception.getMessage()); + assertEquals(1, phaser.getRegisteredParties()); + } + + @Test + void readEntitySwallowsInterruptedCallbacksAndDeregistersReader() throws Exception { + Phaser phaser = new Phaser(1); + + try (MockedConstruction construction = + mockConstruction( + PaginatedEntitiesSource.class, + (mock, context) -> + when(mock.readNextKeyset(isNull())) + .thenReturn((ResultList) mockResult(List.of("table-1"), null, 0)))) { + + int submitted = + reader.readEntity( + "table", + 1, + 10, + phaser, + (entityType, batch, offset) -> { + throw new InterruptedException("stop"); + }); + + assertEquals(1, submitted); + assertEquals(1, phaser.getRegisteredParties()); + assertTrue(Thread.currentThread().isInterrupted()); + Thread.interrupted(); + verify(construction.constructed().get(1)).readNextKeyset(null); + } + } + + @Test + void helperMethodsRespectTimeSeriesAndMinimumReaderRules() { + assertEquals( + List.of(), + ReindexingUtil.getSearchIndexFields(ReportData.ReportDataType.ENTITY_REPORT_DATA.value())); + assertEquals(List.of("*"), ReindexingUtil.getSearchIndexFields("table")); + assertEquals(1, EntityReader.calculateNumberOfReaders(10, 0)); + assertEquals(3, EntityReader.calculateNumberOfReaders(11, 5)); + } + + @Test + void stopAndCloseSetStoppedFlag() { + reader.stop(); + assertTrue(stopped.get()); + + stopped.set(false); + reader.close(); + assertTrue(stopped.get()); + } + + private ResultList mockResult(List data, String after, Integer warningsCount) { + ResultList result = new ResultList<>(); + result.setData(new ArrayList<>(data)); + result.setErrors(null); + result.setWarningsCount(warningsCount); + result.setPaging(new Paging().withAfter(after)); + return result; + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderRetryTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderRetryTest.java new file mode 100644 index 000000000000..8a7fb8710517 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/EntityReaderRetryTest.java @@ -0,0 +1,108 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.service.exception.SearchIndexException; + +@DisplayName("EntityReader Retry Tests") +class EntityReaderRetryTest { + + @Test + @DisplayName("isTransientError detects timeout errors") + void detectsTimeoutErrors() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("Connection timeout while reading entities")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError detects connection errors") + void detectsConnectionErrors() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("java.net.ConnectException: Connection refused")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError detects pool exhaustion") + void detectsPoolExhaustion() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("Pool exhausted - no connections available")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError detects socket timeout") + void detectsSocketTimeout() { + SearchIndexException e = + new SearchIndexException( + new IndexingError().withMessage("java.net.SocketTimeoutException: Read timed out")); + assertTrue(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError returns false for non-transient errors") + void rejectsNonTransientErrors() { + SearchIndexException e = + new SearchIndexException(new IndexingError().withMessage("Entity not found: table.xyz")); + assertFalse(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("isTransientError returns false for null message") + void handleNullMessage() { + SearchIndexException e = new SearchIndexException(new IndexingError()); + assertFalse(EntityReader.isTransientError(e)); + } + + @Test + @DisplayName("EntityReader constructor accepts custom retry configuration") + void customRetryConfiguration() { + java.util.concurrent.ExecutorService executor = + java.util.concurrent.Executors.newSingleThreadExecutor(); + java.util.concurrent.atomic.AtomicBoolean stopped = + new java.util.concurrent.atomic.AtomicBoolean(false); + EntityReader reader = new EntityReader(executor, stopped, 5, 1000); + assertNotNull(reader); + executor.shutdown(); + } + + @Test + @DisplayName("EntityReader default constructor uses default retry values") + void defaultRetryConfiguration() { + java.util.concurrent.ExecutorService executor = + java.util.concurrent.Executors.newSingleThreadExecutor(); + java.util.concurrent.atomic.AtomicBoolean stopped = + new java.util.concurrent.atomic.AtomicBoolean(false); + EntityReader reader = new EntityReader(executor, stopped); + assertNotNull(reader); + executor.shutdown(); + } + + @Test + @DisplayName("VectorCompletionResult.success creates completed result") + void vectorCompletionSuccess() { + VectorCompletionResult result = VectorCompletionResult.success(150); + assertTrue(result.completed()); + assertEquals(0, result.pendingTaskCount()); + assertEquals(150, result.waitedMillis()); + } + + @Test + @DisplayName("VectorCompletionResult.timeout creates timeout result") + void vectorCompletionTimeout() { + VectorCompletionResult result = VectorCompletionResult.timeout(5, 30000); + assertFalse(result.completed()); + assertEquals(5, result.pendingTaskCount()); + assertEquals(30000, result.waitedMillis()); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipelineTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipelineTest.java new file mode 100644 index 000000000000..c53c7119589e --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/IndexingPipelineTest.java @@ -0,0 +1,473 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Phaser; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.system.EntityStats; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.Entity; +import org.openmetadata.service.jdbi3.EntityDAO; +import org.openmetadata.service.jdbi3.EntityRepository; +import org.openmetadata.service.jdbi3.EntityTimeSeriesDAO; +import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; +import org.openmetadata.service.jdbi3.ListFilter; +import org.openmetadata.service.search.EntityReindexContext; +import org.openmetadata.service.search.RecreateIndexHandler; +import org.openmetadata.service.search.ReindexContext; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.util.FullyQualifiedName; + +@SuppressWarnings({"rawtypes", "unchecked"}) +class IndexingPipelineTest { + + private SearchRepository searchRepository; + private IndexingPipeline pipeline; + + @BeforeEach + void setUp() { + searchRepository = mock(SearchRepository.class); + pipeline = new IndexingPipeline(searchRepository); + } + + @AfterEach + void tearDown() { + pipeline.close(); + } + + @Test + void executeProcessesEntitiesUsingComputedTotalsAndCompletes() throws Exception { + BulkSink sink = mock(BulkSink.class); + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ReindexingJobContext context = mockJobContext(); + EntityRepository entityRepository = mock(EntityRepository.class); + EntityDAO entityDao = mock(EntityDAO.class); + EntityInterface entityA = mock(EntityInterface.class); + EntityInterface entityB = mock(EntityInterface.class); + ResultList batch = new ResultList<>(List.of(entityA, entityB), null, null, 0); + + when(entityRepository.getDao()).thenReturn(entityDao); + when(entityDao.listCount(any(ListFilter.class))).thenReturn(2); + when(sink.getPendingVectorTaskCount()).thenReturn(0); + when(sink.getStats()).thenReturn(new StepStats().withTotalRecords(2).withSuccessRecords(2)); + when(sink.getProcessStats()) + .thenReturn(new StepStats().withTotalRecords(2).withSuccessRecords(2)); + + pipeline.addListener(listener); + + try (MockedStatic entityMock = mockStatic(Entity.class); + MockedConstruction ignored = + mockConstruction( + EntityReader.class, + (reader, context1) -> + doAnswer( + invocation -> { + String entityType = invocation.getArgument(0); + int totalRecords = invocation.getArgument(1); + EntityReader.BatchCallback callback = invocation.getArgument(4); + assertEquals(Entity.TABLE, entityType); + assertEquals(2, totalRecords); + callback.onBatchRead(entityType, batch, 0); + return 1; + }) + .when(reader) + .readEntity( + any(String.class), + anyInt(), + anyInt(), + any(Phaser.class), + any(EntityReader.BatchCallback.class), + any(), + any()))) { + entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); + + ExecutionResult result = + pipeline.execute( + ReindexingConfiguration.builder() + .entities(Set.of(Entity.TABLE)) + .batchSize(2) + .consumerThreads(1) + .producerThreads(1) + .build(), + context, + Set.of(Entity.TABLE), + sink, + null, + null); + + assertEquals(ExecutionResult.Status.COMPLETED, result.status()); + assertEquals(2, result.finalStats().getJobStats().getTotalRecords()); + assertEquals(2, result.finalStats().getJobStats().getSuccessRecords()); + + ArgumentCaptor dataCaptor = ArgumentCaptor.forClass(List.class); + ArgumentCaptor contextCaptor = ArgumentCaptor.forClass(Map.class); + verify(sink).write(dataCaptor.capture(), contextCaptor.capture()); + assertEquals(2, dataCaptor.getValue().size()); + assertEquals(Entity.TABLE, contextCaptor.getValue().get("entityType")); + assertEquals(Boolean.FALSE, contextCaptor.getValue().get("recreateIndex")); + + verify(listener).onJobStarted(context); + verify(listener).onEntityTypeStarted(Entity.TABLE, 2); + verify(listener).onProgressUpdate(any(Stats.class), isNull()); + verify(listener).onJobCompleted(any(Stats.class), anyLong()); + } + } + + @Test + void executeMarksCompletedWithErrorsWhenSinkWriteFails() throws Exception { + BulkSink sink = mock(BulkSink.class); + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ReindexingJobContext context = mockJobContext(); + EntityRepository entityRepository = mock(EntityRepository.class); + EntityDAO entityDao = mock(EntityDAO.class); + EntityInterface entity = mock(EntityInterface.class); + ResultList batch = new ResultList<>(List.of(entity), null, null, 0); + + when(entityRepository.getDao()).thenReturn(entityDao); + when(entityDao.listCount(any(ListFilter.class))).thenReturn(1); + when(sink.getPendingVectorTaskCount()).thenReturn(0); + when(sink.getStats()).thenReturn(new StepStats().withTotalRecords(0).withSuccessRecords(0)); + when(sink.getProcessStats()) + .thenReturn(new StepStats().withTotalRecords(0).withSuccessRecords(0)); + pipeline.addListener(listener); + + try (MockedStatic entityMock = mockStatic(Entity.class); + MockedConstruction ignored = + mockConstruction( + EntityReader.class, + (reader, context1) -> + doAnswer( + invocation -> { + EntityReader.BatchCallback callback = invocation.getArgument(4); + callback.onBatchRead(Entity.TABLE, batch, 0); + return 1; + }) + .when(reader) + .readEntity( + any(String.class), + anyInt(), + anyInt(), + any(Phaser.class), + any(EntityReader.BatchCallback.class), + any(), + any()))) { + entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); + doAnswer( + invocation -> { + throw new IllegalStateException("sink boom"); + }) + .when(sink) + .write(any(List.class), any(Map.class)); + + ExecutionResult result = + pipeline.execute( + ReindexingConfiguration.builder() + .entities(Set.of(Entity.TABLE)) + .batchSize(1) + .consumerThreads(1) + .producerThreads(1) + .build(), + context, + Set.of(Entity.TABLE), + sink, + null, + null); + + assertEquals(ExecutionResult.Status.COMPLETED_WITH_ERRORS, result.status()); + assertEquals(1, result.finalStats().getJobStats().getTotalRecords()); + assertEquals(0, result.finalStats().getJobStats().getSuccessRecords()); + + ArgumentCaptor errorCaptor = ArgumentCaptor.forClass(IndexingError.class); + verify(listener).onError(eq(Entity.TABLE), errorCaptor.capture(), any(Stats.class)); + assertEquals(IndexingError.ErrorSource.SINK, errorCaptor.getValue().getErrorSource()); + assertEquals("sink boom", errorCaptor.getValue().getMessage()); + verify(listener).onJobCompletedWithErrors(any(Stats.class), anyLong()); + } + } + + @Test + void initializeStatsUsesRepositoryTotalsForRegularAndTimeSeriesEntities() throws Exception { + EntityRepository entityRepository = mock(EntityRepository.class); + EntityDAO entityDao = mock(EntityDAO.class); + EntityTimeSeriesRepository timeSeriesRepository = mock(EntityTimeSeriesRepository.class); + EntityTimeSeriesDAO timeSeriesDao = mock(EntityTimeSeriesDAO.class); + String reportType = ReportData.ReportDataType.ENTITY_REPORT_DATA.value(); + + when(entityRepository.getDao()).thenReturn(entityDao); + when(entityDao.listCount(any(ListFilter.class))).thenReturn(7); + when(timeSeriesRepository.getTimeSeriesDao()).thenReturn(timeSeriesDao); + when(timeSeriesDao.listCount(any(ListFilter.class), anyLong(), anyLong(), eq(false))) + .thenReturn(3); + when(searchRepository.getDataInsightReports()).thenReturn(List.of(reportType)); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); + entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); + entityMock + .when(() -> Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA)) + .thenReturn(timeSeriesRepository); + + Stats stats = + (Stats) + invokePrivate( + "initializeStats", + new Class[] {ReindexingConfiguration.class, Set.class}, + ReindexingConfiguration.builder() + .timeSeriesEntityDays(Map.of(reportType, 1)) + .build(), + Set.of(Entity.TABLE, reportType)); + + assertEquals(10, stats.getJobStats().getTotalRecords()); + assertEquals(10, stats.getReaderStats().getTotalRecords()); + assertEquals( + 7, stats.getEntityStats().getAdditionalProperties().get(Entity.TABLE).getTotalRecords()); + assertEquals( + 3, stats.getEntityStats().getAdditionalProperties().get(reportType).getTotalRecords()); + + ArgumentCaptor filterCaptor = ArgumentCaptor.forClass(ListFilter.class); + verify(timeSeriesDao).listCount(filterCaptor.capture(), anyLong(), anyLong(), eq(false)); + assertEquals( + FullyQualifiedName.buildHash(reportType), + filterCaptor.getValue().getQueryParams().get("entityFQNHash")); + } + } + + @Test + void getEntityTotalUsesEntitySpecificTimeSeriesRepositoryWithoutTimeWindow() throws Exception { + EntityTimeSeriesRepository timeSeriesRepository = mock(EntityTimeSeriesRepository.class); + EntityTimeSeriesDAO timeSeriesDao = mock(EntityTimeSeriesDAO.class); + String entityType = ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(); + + when(timeSeriesRepository.getTimeSeriesDao()).thenReturn(timeSeriesDao); + when(timeSeriesDao.listCount(any(ListFilter.class))).thenReturn(5); + when(searchRepository.getDataInsightReports()).thenReturn(List.of()); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock + .when(() -> Entity.getEntityTimeSeriesRepository(entityType)) + .thenReturn(timeSeriesRepository); + entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); + + int total = + (int) + invokePrivate( + "getEntityTotal", + new Class[] {String.class, ReindexingConfiguration.class}, + entityType, + null); + + assertEquals(5, total); + verify(timeSeriesDao).listCount(any(ListFilter.class)); + } + } + + @Test + void getEntityTotalReturnsZeroWhenTimeSeriesRepositoryCountFails() throws Exception { + EntityTimeSeriesRepository timeSeriesRepository = mock(EntityTimeSeriesRepository.class); + EntityTimeSeriesDAO timeSeriesDao = mock(EntityTimeSeriesDAO.class); + String entityType = ReportData.ReportDataType.WEB_ANALYTIC_USER_ACTIVITY_REPORT_DATA.value(); + + when(timeSeriesRepository.getTimeSeriesDao()).thenReturn(timeSeriesDao); + when(timeSeriesDao.listCount(any(ListFilter.class))) + .thenThrow(new IllegalStateException("boom")); + when(searchRepository.getDataInsightReports()).thenReturn(List.of()); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock + .when(() -> Entity.getEntityTimeSeriesRepository(entityType)) + .thenReturn(timeSeriesRepository); + entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); + + int total = + (int) + invokePrivate( + "getEntityTotal", + new Class[] {String.class, ReindexingConfiguration.class}, + entityType, + null); + + assertEquals(0, total); + } + } + + @Test + void createContextDataAndFinalizeReindexUseRecreateMetadata() throws Exception { + ReindexContext recreateContext = new ReindexContext(); + recreateContext.add( + Entity.TABLE, + "table-canonical", + "table-original", + "table-staged", + Set.of("table-alias"), + "table-canonical-alias", + List.of("table-parent")); + recreateContext.add( + Entity.USER, + "user-canonical", + "user-original", + "user-staged", + Set.of("user-alias"), + "user-canonical-alias", + List.of("user-parent")); + RecreateIndexHandler handler = mock(RecreateIndexHandler.class); + + setField("recreateContext", recreateContext); + setField("recreateIndexHandler", handler); + getPromotedEntities().add(Entity.TABLE); + + Map contextData = + (Map) + invokePrivate("createContextData", new Class[] {String.class}, Entity.TABLE); + + assertEquals(Entity.TABLE, contextData.get("entityType")); + assertEquals(Boolean.TRUE, contextData.get("recreateIndex")); + assertSame(recreateContext, contextData.get("recreateContext")); + assertEquals("table-staged", contextData.get("targetIndex")); + + invokePrivate("finalizeReindex", new Class[0]); + + ArgumentCaptor contextCaptor = + ArgumentCaptor.forClass(EntityReindexContext.class); + verify(handler).finalizeReindex(contextCaptor.capture(), eq(true)); + assertEquals(Entity.USER, contextCaptor.getValue().getEntityType()); + assertEquals("user-canonical", contextCaptor.getValue().getCanonicalIndex()); + assertEquals("user-original", contextCaptor.getValue().getOriginalIndex()); + assertEquals("user-staged", contextCaptor.getValue().getStagedIndex()); + assertTrue(contextCaptor.getValue().getExistingAliases().contains("user-alias")); + assertTrue(contextCaptor.getValue().getParentAliases().contains("user-parent")); + assertNull(getField("recreateContext")); + assertTrue(getPromotedEntities().isEmpty()); + } + + @Test + void buildResultReturnsStoppedAndNotifiesListeners() throws Exception { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + pipeline.addListener(listener); + pipeline.getStats().set(createStats("table", 2)); + getStoppedFlag().set(true); + + ExecutionResult result = + (ExecutionResult) + invokePrivate( + "buildResult", new Class[] {long.class}, System.currentTimeMillis() - 1000); + + assertEquals(ExecutionResult.Status.STOPPED, result.status()); + verify(listener).onJobStopped(any(Stats.class)); + } + + @Test + void stopFlushesSinkStopsReaderAndShutsExecutorsDown() throws Exception { + BulkSink sink = mock(BulkSink.class); + EntityReader reader = mock(EntityReader.class); + LinkedBlockingQueue queue = new LinkedBlockingQueue(); + ExecutorService producerExecutor = Executors.newSingleThreadExecutor(); + ExecutorService jobExecutor = Executors.newSingleThreadExecutor(); + ExecutorService consumerExecutor = Executors.newSingleThreadExecutor(); + + queue.offer("pending-task"); + when(sink.getActiveBulkRequestCount()).thenReturn(2); + when(sink.flushAndAwait(10)).thenReturn(true); + + setField("searchIndexSink", sink); + setField("entityReader", reader); + setField("taskQueue", queue); + setField("producerExecutor", producerExecutor); + setField("jobExecutor", jobExecutor); + setField("consumerExecutor", consumerExecutor); + + pipeline.stop(); + + assertTrue(getStoppedFlag().get()); + assertFalse(queue.isEmpty()); + verify(reader).stop(); + verify(sink).flushAndAwait(10); + assertTrue(producerExecutor.isShutdown()); + assertTrue(jobExecutor.isShutdown()); + assertTrue(consumerExecutor.isShutdown()); + } + + private ReindexingJobContext mockJobContext() { + ReindexingJobContext context = mock(ReindexingJobContext.class); + when(context.getJobId()).thenReturn(UUID.fromString("00000000-0000-0000-0000-000000000041")); + when(context.getJobName()).thenReturn("job"); + when(context.getStartTime()).thenReturn(System.currentTimeMillis()); + when(context.isDistributed()).thenReturn(false); + when(context.getSource()).thenReturn("TEST"); + return context; + } + + private Stats createStats(String entityType, int totalRecords) { + Stats stats = new Stats(); + EntityStats entityStats = new EntityStats(); + entityStats.withAdditionalProperty( + entityType, new StepStats().withTotalRecords(totalRecords).withSuccessRecords(0)); + stats.setEntityStats(entityStats); + stats.setJobStats(new StepStats().withTotalRecords(totalRecords).withSuccessRecords(0)); + stats.setReaderStats(new StepStats().withTotalRecords(totalRecords).withSuccessRecords(0)); + stats.setSinkStats(new StepStats().withTotalRecords(0).withSuccessRecords(0)); + stats.setProcessStats(new StepStats().withTotalRecords(0).withSuccessRecords(0)); + return stats; + } + + private AtomicBoolean getStoppedFlag() throws Exception { + return (AtomicBoolean) getField("stopped"); + } + + private Set getPromotedEntities() throws Exception { + return (Set) getField("promotedEntities"); + } + + private Object invokePrivate(String methodName, Class[] parameterTypes, Object... args) + throws Exception { + Method method = IndexingPipeline.class.getDeclaredMethod(methodName, parameterTypes); + method.setAccessible(true); + return method.invoke(pipeline, args); + } + + private void setField(String fieldName, Object value) throws Exception { + Field field = IndexingPipeline.class.getDeclaredField(fieldName); + field.setAccessible(true); + field.set(pipeline, value); + } + + private Object getField(String fieldName) throws Exception { + Field field = IndexingPipeline.class.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(pipeline); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContextTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContextTest.java index 7c1b5707ed63..ed579ccc2c61 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContextTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzJobContextTest.java @@ -1,6 +1,7 @@ package org.openmetadata.service.apps.bundles.searchIndex; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -29,23 +30,25 @@ void quartzJobContextUsesQuartzAndAppMetadata() { when(app.getId()).thenReturn(appId); long before = System.currentTimeMillis(); - QuartzJobContext context = new QuartzJobContext(quartzContext, app); + QuartzJobContext context = new QuartzJobContext(quartzContext, app, true); long after = System.currentTimeMillis(); assertEquals(appId, context.getJobId()); assertEquals("reindex-job", context.getJobName()); assertEquals(appId, context.getAppId()); assertTrue(context.getStartTime() >= before && context.getStartTime() <= after); + assertTrue(context.isDistributed()); assertEquals("QUARTZ", context.getSource()); } @Test void quartzJobContextFallsBackWhenQuartzContextOrAppIsMissing() { - QuartzJobContext context = new QuartzJobContext(null, null); + QuartzJobContext context = new QuartzJobContext(null, null, false); assertNotNull(context.getJobId()); assertEquals("unknown", context.getJobName()); assertNull(context.getAppId()); + assertFalse(context.isDistributed()); assertEquals("QUARTZ", context.getSource()); } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContextTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContextTest.java index 5becc78c670b..f9f977f0e0c6 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContextTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/QuartzOrchestratorContextTest.java @@ -68,7 +68,7 @@ void orchestratorContextDelegatesQuartzStorageAndFactoryMethods() { QuartzProgressListener.class, context.createProgressListener( new EventPublisherJob().withEntities(java.util.Set.of("table")))); - assertInstanceOf(QuartzJobContext.class, context.createReindexingContext()); + assertInstanceOf(QuartzJobContext.class, context.createReindexingContext(true)); } @Test diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfigurationTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfigurationTest.java deleted file mode 100644 index 9f60aaa198c7..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingConfigurationTest.java +++ /dev/null @@ -1,27 +0,0 @@ -package org.openmetadata.service.apps.bundles.searchIndex; - -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.Set; -import org.junit.jupiter.api.Test; -import org.openmetadata.service.Entity; - -class ReindexingConfigurationTest { - - @Test - void isSmartReindexingReturnsFalseForAllEntities() { - ReindexingConfiguration config = - ReindexingConfiguration.builder().entities(Set.of(SearchIndexEntityTypes.ALL)).build(); - - assertFalse(config.isSmartReindexing()); - } - - @Test - void isSmartReindexingReturnsTrueForSmallEntitySubsets() { - ReindexingConfiguration config = - ReindexingConfiguration.builder().entities(Set.of(Entity.TABLE)).build(); - - assertTrue(config.isSmartReindexing()); - } -} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestratorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestratorTest.java index 495f32c7b98c..5dd15adc4a9f 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestratorTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/ReindexingOrchestratorTest.java @@ -31,7 +31,6 @@ import java.util.UUID; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; import org.mockito.MockedConstruction; import org.mockito.MockedStatic; import org.openmetadata.schema.api.configuration.OpenMetadataBaseUrlConfiguration; @@ -93,9 +92,12 @@ void setUp() { } @Test - void runPreservesResultMetadataInSuccessContext() { + void runSingleServerPreservesResultMetadataInSuccessContext() { EventPublisherJob jobData = - new EventPublisherJob().withEntities(Set.of(Entity.TABLE)).withBatchSize(25); + new EventPublisherJob() + .withEntities(Set.of(Entity.TABLE)) + .withBatchSize(25) + .withUseDistributedIndexing(false); ReindexingProgressListener progressListener = mock(ReindexingProgressListener.class); ReindexingJobContext jobContext = mock(ReindexingJobContext.class); EntityRepository entityRepository = mock(EntityRepository.class); @@ -104,7 +106,7 @@ void runPreservesResultMetadataInSuccessContext() { when(context.getJobName()).thenReturn("scheduled"); when(context.createProgressListener(jobData)).thenReturn(progressListener); - when(context.createReindexingContext()).thenReturn(jobContext); + when(context.createReindexingContext(false)).thenReturn(jobContext); when(searchIndexFailureDAO.countByJobId(appRunRecord.getAppId().toString())).thenReturn(0); when(entityRepository.getDao()).thenReturn(entityDao); when(entityDao.listCount(any())).thenReturn(5); @@ -113,9 +115,9 @@ void runPreservesResultMetadataInSuccessContext() { MockedStatic metricsMock = mockStatic(ReindexingMetrics.class); MockedStatic websocketMock = mockStatic(WebSocketManager.class); MockedConstruction cleanerConstruction = mockOrphanCleaner(); - MockedConstruction strategyConstruction = + MockedConstruction strategyConstruction = mockConstruction( - DistributedIndexingStrategy.class, + SingleServerIndexingStrategy.class, (strategy, context1) -> { when(strategy.execute(any(), any())) .thenReturn( @@ -137,7 +139,7 @@ void runPreservesResultMetadataInSuccessContext() { orchestrator.run(jobData); - DistributedIndexingStrategy strategy = strategyConstruction.constructed().getFirst(); + SingleServerIndexingStrategy strategy = strategyConstruction.constructed().getFirst(); verify(strategy, times(2)).addListener(any(ReindexingProgressListener.class)); verify(strategy).execute(any(ReindexingConfiguration.class), eq(jobContext)); verify(context).storeRunStats(stats); @@ -163,8 +165,8 @@ void runLoadsOnDemandConfigAndCompletesWithoutBuildingStrategy() { try (MockedStatic metricsMock = mockStatic(ReindexingMetrics.class); MockedStatic websocketMock = mockStatic(WebSocketManager.class); MockedConstruction ignoredCleaner = mockOrphanCleaner(); - MockedConstruction ignoredStrategy = - mockConstruction(DistributedIndexingStrategy.class)) { + MockedConstruction ignoredStrategy = + mockConstruction(SingleServerIndexingStrategy.class)) { metricsMock.when(ReindexingMetrics::getInstance).thenReturn(null); websocketMock.when(WebSocketManager::getInstance).thenReturn(null); @@ -179,38 +181,6 @@ void runLoadsOnDemandConfigAndCompletesWithoutBuildingStrategy() { } } - @Test - void runRemovesLegacyModeOptionsFromOnDemandAndRunRecordConfig() { - EventPublisherJob jobData = new EventPublisherJob().withEntities(Set.of()); - Map legacyConfig = JsonUtils.convertValue(jobData, Map.class); - legacyConfig.put("recreateIndex", true); - legacyConfig.put("useDistributedIndexing", false); - appRunRecord.setConfig(new HashMap<>(legacyConfig)); - - when(context.getJobName()).thenReturn(ON_DEMAND_JOB); - when(context.getAppConfigJson()).thenReturn(JsonUtils.pojoToJson(legacyConfig)); - when(searchIndexFailureDAO.countByJobId(appRunRecord.getAppId().toString())).thenReturn(0); - - try (MockedStatic metricsMock = mockStatic(ReindexingMetrics.class); - MockedStatic websocketMock = mockStatic(WebSocketManager.class); - MockedConstruction ignoredCleaner = mockOrphanCleaner(); - MockedConstruction ignoredStrategy = - mockConstruction(DistributedIndexingStrategy.class)) { - metricsMock.when(ReindexingMetrics::getInstance).thenReturn(null); - websocketMock.when(WebSocketManager::getInstance).thenReturn(null); - - orchestrator.run(null); - - ArgumentCaptor configCaptor = ArgumentCaptor.forClass(Map.class); - verify(context).updateAppConfiguration(configCaptor.capture()); - assertFalse(configCaptor.getValue().containsKey("recreateIndex")); - assertFalse(configCaptor.getValue().containsKey("useDistributedIndexing")); - assertFalse(appRunRecord.getConfig().containsKey("recreateIndex")); - assertFalse(appRunRecord.getConfig().containsKey("useDistributedIndexing")); - assertTrue(ignoredStrategy.constructed().isEmpty()); - } - } - @Test void runContinuesWhenHybridPipelinePreflightFails() { EventPublisherJob jobData = new EventPublisherJob().withEntities(Set.of()); @@ -225,8 +195,8 @@ void runContinuesWhenHybridPipelinePreflightFails() { try (MockedStatic metricsMock = mockStatic(ReindexingMetrics.class); MockedStatic websocketMock = mockStatic(WebSocketManager.class); MockedConstruction ignoredCleaner = mockOrphanCleaner(); - MockedConstruction ignoredStrategy = - mockConstruction(DistributedIndexingStrategy.class)) { + MockedConstruction ignoredStrategy = + mockConstruction(SingleServerIndexingStrategy.class)) { metricsMock.when(ReindexingMetrics::getInstance).thenReturn(null); websocketMock.when(WebSocketManager::getInstance).thenReturn(null); @@ -241,7 +211,10 @@ void runContinuesWhenHybridPipelinePreflightFails() { @Test void runMarksJobFailedAndCapturesStrategyStatsOnExecutionException() { - EventPublisherJob jobData = new EventPublisherJob().withEntities(Set.of(Entity.TABLE)); + EventPublisherJob jobData = + new EventPublisherJob() + .withEntities(Set.of(Entity.TABLE)) + .withUseDistributedIndexing(false); ReindexingProgressListener progressListener = mock(ReindexingProgressListener.class); ReindexingJobContext jobContext = mock(ReindexingJobContext.class); EntityRepository entityRepository = mock(EntityRepository.class); @@ -250,7 +223,7 @@ void runMarksJobFailedAndCapturesStrategyStatsOnExecutionException() { when(context.getJobName()).thenReturn("scheduled"); when(context.createProgressListener(jobData)).thenReturn(progressListener); - when(context.createReindexingContext()).thenReturn(jobContext); + when(context.createReindexingContext(false)).thenReturn(jobContext); when(searchIndexFailureDAO.countByJobId(appRunRecord.getAppId().toString())).thenReturn(0); when(entityRepository.getDao()).thenReturn(entityDao); when(entityDao.listCount(any())).thenReturn(3); @@ -259,9 +232,9 @@ void runMarksJobFailedAndCapturesStrategyStatsOnExecutionException() { MockedStatic metricsMock = mockStatic(ReindexingMetrics.class); MockedStatic websocketMock = mockStatic(WebSocketManager.class); MockedConstruction ignoredCleaner = mockOrphanCleaner(); - MockedConstruction ignoredStrategy = + MockedConstruction ignoredStrategy = mockConstruction( - DistributedIndexingStrategy.class, + SingleServerIndexingStrategy.class, (strategy, context1) -> { when(strategy.execute(any(), any())).thenThrow(new RuntimeException("boom")); when(strategy.getStats()).thenReturn(Optional.of(stats)); @@ -284,7 +257,7 @@ void runMarksJobFailedAndCapturesStrategyStatsOnExecutionException() { @Test void stopStopsActiveStrategyAndPushesStoppedStatus() throws Exception { - DistributedIndexingStrategy strategy = mock(DistributedIndexingStrategy.class); + IndexingStrategy strategy = mock(IndexingStrategy.class); EventPublisherJob jobData = new EventPublisherJob() .withEntities(Set.of(Entity.TABLE)) @@ -339,7 +312,8 @@ void runAddsSlackListenerUsingInstanceUrlFromSettings() { new EventPublisherJob() .withEntities(Set.of(Entity.TABLE)) .withSlackBotToken("token") - .withSlackChannel("#alerts"); + .withSlackChannel("#alerts") + .withUseDistributedIndexing(false); ReindexingProgressListener progressListener = mock(ReindexingProgressListener.class); ReindexingJobContext jobContext = mock(ReindexingJobContext.class); EntityRepository entityRepository = mock(EntityRepository.class); @@ -349,7 +323,7 @@ void runAddsSlackListenerUsingInstanceUrlFromSettings() { when(context.getJobName()).thenReturn("scheduled"); when(context.createProgressListener(jobData)).thenReturn(progressListener); - when(context.createReindexingContext()).thenReturn(jobContext); + when(context.createReindexingContext(false)).thenReturn(jobContext); when(searchIndexFailureDAO.countByJobId(appRunRecord.getAppId().toString())).thenReturn(0); when(entityRepository.getDao()).thenReturn(entityDao); when(entityDao.listCount(any())).thenReturn(2); @@ -363,9 +337,9 @@ void runAddsSlackListenerUsingInstanceUrlFromSettings() { MockedStatic metricsMock = mockStatic(ReindexingMetrics.class); MockedStatic websocketMock = mockStatic(WebSocketManager.class); MockedConstruction ignoredCleaner = mockOrphanCleaner(); - MockedConstruction strategyConstruction = + MockedConstruction strategyConstruction = mockConstruction( - DistributedIndexingStrategy.class, + SingleServerIndexingStrategy.class, (strategy, context1) -> when(strategy.execute(any(), any())) .thenReturn( @@ -386,7 +360,7 @@ void runAddsSlackListenerUsingInstanceUrlFromSettings() { orchestrator.run(jobData); - DistributedIndexingStrategy strategy = strategyConstruction.constructed().getFirst(); + SingleServerIndexingStrategy strategy = strategyConstruction.constructed().getFirst(); verify(strategy, times(3)).addListener(any(ReindexingProgressListener.class)); verify(context, never()).updateAppConfiguration(any(Map.class)); } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppConfigSanitizerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppConfigSanitizerTest.java deleted file mode 100644 index a8bbe66e579f..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexAppConfigSanitizerTest.java +++ /dev/null @@ -1,44 +0,0 @@ -package org.openmetadata.service.apps.bundles.searchIndex; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNotSame; -import static org.junit.jupiter.api.Assertions.assertNull; - -import java.util.LinkedHashMap; -import java.util.Map; -import org.junit.jupiter.api.Test; - -class SearchIndexAppConfigSanitizerTest { - - @Test - void copyWithoutRemovedOptionsReturnsNullForNullConfig() { - assertNull(SearchIndexAppConfigSanitizer.copyWithoutRemovedOptions(null)); - } - - @Test - void copyWithoutRemovedOptionsReturnsDefensiveCopyForEmptyConfig() { - Map config = new LinkedHashMap<>(); - - Map sanitized = SearchIndexAppConfigSanitizer.copyWithoutRemovedOptions(config); - - assertNotSame(config, sanitized); - assertEquals(config, sanitized); - } - - @Test - void copyWithoutRemovedOptionsRemovesDeprecatedDistributedOptions() { - Map config = new LinkedHashMap<>(); - config.put("batchSize", 100); - config.put("recreateIndex", true); - config.put("useDistributedIndexing", true); - - Map sanitized = SearchIndexAppConfigSanitizer.copyWithoutRemovedOptions(config); - - assertNotSame(config, sanitized); - assertEquals(100, sanitized.get("batchSize")); - assertFalse(sanitized.containsKey("recreateIndex")); - assertFalse(sanitized.containsKey("useDistributedIndexing")); - assertEquals(3, config.size()); - } -} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexEndToEndTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexEndToEndTest.java new file mode 100644 index 000000000000..2c5c7ea6c0be --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexEndToEndTest.java @@ -0,0 +1,416 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.junit.jupiter.MockitoExtension; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.entity.app.App; +import org.openmetadata.schema.entity.app.AppRunRecord; +import org.openmetadata.schema.system.EntityError; +import org.openmetadata.schema.system.EventPublisherJob; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.schema.utils.JsonUtils; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.exception.SearchIndexException; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.socket.WebSocketManager; +import org.quartz.JobDataMap; +import org.quartz.JobDetail; +import org.quartz.JobExecutionContext; + +/** + * End-to-end test that verifies the complete fix for: + * 1. Error propagation from ElasticSearchIndexSink to SearchIndexExecutor + * 2. Real-time WebSocket updates for metrics and errors + * 3. Proper job completion status + * 4. Field limit error handling specifically + */ +@ExtendWith(MockitoExtension.class) +@Slf4j +public class SearchIndexEndToEndTest { + + @Mock private CollectionDAO collectionDAO; + @Mock private SearchRepository searchRepository; + @Mock private BulkSink mockSink; + @Mock private JobExecutionContext jobExecutionContext; + @Mock private JobDetail jobDetail; + @Mock private JobDataMap jobDataMap; + @Mock private WebSocketManager webSocketManager; + @Mock private org.quartz.Scheduler scheduler; + @Mock private org.quartz.ListenerManager listenerManager; + @Mock private org.openmetadata.service.apps.scheduler.OmAppJobListener jobListener; + @Mock private AppRunRecord appRunRecord; + + private SearchIndexApp searchIndexApp; + private SearchIndexExecutor searchIndexExecutor; + private final ObjectMapper objectMapper = new ObjectMapper(); + private final List webSocketMessages = + Collections.synchronizedList(new ArrayList<>()); + private MockedStatic webSocketManagerMock; + + private static class WebSocketMessage { + String channel; + String content; + long timestamp; + + WebSocketMessage(String channel, String content) { + this.channel = channel; + this.content = content; + this.timestamp = System.currentTimeMillis(); + } + } + + @BeforeEach + void setUp() { + searchIndexApp = new SearchIndexApp(collectionDAO, searchRepository); + searchIndexExecutor = new SearchIndexExecutor(collectionDAO, searchRepository); + lenient().when(jobExecutionContext.getJobDetail()).thenReturn(jobDetail); + lenient().when(jobDetail.getJobDataMap()).thenReturn(jobDataMap); + lenient().when(jobDataMap.get("triggerType")).thenReturn("MANUAL"); + + try { + lenient().when(jobExecutionContext.getScheduler()).thenReturn(scheduler); + lenient().when(scheduler.getListenerManager()).thenReturn(listenerManager); + lenient().when(listenerManager.getJobListener(anyString())).thenReturn(jobListener); + lenient().when(jobListener.getAppRunRecordForJob(any())).thenReturn(appRunRecord); + lenient().when(appRunRecord.getStatus()).thenReturn(AppRunRecord.Status.RUNNING); + } catch (Exception e) { + // Ignore mocking exceptions in test setup + } + + webSocketManagerMock = mockStatic(WebSocketManager.class); + webSocketManagerMock.when(WebSocketManager::getInstance).thenReturn(webSocketManager); + + lenient() + .doAnswer( + invocation -> { + String channel = invocation.getArgument(0); + String content = invocation.getArgument(1); + webSocketMessages.add(new WebSocketMessage(channel, content)); + LOG.debug( + "WebSocket message captured - Channel: {}, Content length: {}", + channel, + content.length()); + return null; + }) + .when(webSocketManager) + .broadCastMessageToAll(anyString(), anyString()); + } + + @AfterEach + void tearDown() { + if (webSocketManagerMock != null) { + webSocketManagerMock.close(); + } + if (searchIndexExecutor != null) { + searchIndexExecutor.close(); + } + } + + @Test + void testCompleteFieldLimitErrorFlow() throws Exception { + EventPublisherJob jobData = + new EventPublisherJob() + .withEntities(Set.of("table")) + .withBatchSize(5) + .withPayLoadSize(1000000L) + .withMaxConcurrentRequests(10) + .withMaxRetries(3) + .withInitialBackoff(1000) + .withMaxBackoff(10000) + .withProducerThreads(1) + .withConsumerThreads(1) + .withQueueSize(50) + .withRecreateIndex(false) + .withStats(new Stats()); + + App testApp = + new App() + .withName("SearchIndexingApplication") + .withAppConfiguration(JsonUtils.convertValue(jobData, Object.class)); + + ReindexingConfiguration config = ReindexingConfiguration.from(jobData); + + try { + java.lang.reflect.Field configField = SearchIndexExecutor.class.getDeclaredField("config"); + configField.setAccessible(true); + configField.set(searchIndexExecutor, config); + + java.lang.reflect.Field sinkField = + SearchIndexExecutor.class.getDeclaredField("searchIndexSink"); + sinkField.setAccessible(true); + sinkField.set(searchIndexExecutor, mockSink); + + Stats initialStats = searchIndexExecutor.initializeTotalRecords(jobData.getEntities()); + searchIndexExecutor.getStats().set(initialStats); + } catch (Exception e) { + throw new RuntimeException("Failed to set fields via reflection", e); + } + webSocketMessages.clear(); + + List entities = new ArrayList<>(); + for (int i = 0; i < 10; i++) { + EntityInterface entity = mock(EntityInterface.class); + lenient().when(entity.getId()).thenReturn(UUID.randomUUID()); + entities.add(entity); + } + + List fieldLimitErrors = + Arrays.asList( + new EntityError() + .withMessage( + "Elasticsearch exception [type=document_parsing_exception, reason=[1:6347] failed to parse: Limit of total fields [250] has been exceeded while adding new fields [3]]") + .withEntity("table_entity_1"), + new EntityError() + .withMessage( + "Elasticsearch exception [type=document_parsing_exception, reason=[1:3302] failed to parse: Limit of total fields [250] has been exceeded while adding new fields [1]]") + .withEntity("table_entity_2"), + new EntityError() + .withMessage( + "Elasticsearch exception [type=document_parsing_exception, reason=[1:1651] failed to parse: Limit of total fields [250] has been exceeded while adding new fields [1]]") + .withEntity("table_entity_3")); + + IndexingError sinkError = + new IndexingError() + .withErrorSource(IndexingError.ErrorSource.SINK) + .withSubmittedCount(10) + .withSuccessCount(7) + .withFailedCount(3) + .withMessage("Issues in Sink to Elasticsearch: Field limit exceeded") + .withFailedEntities(fieldLimitErrors); + + SearchIndexException sinkException = new SearchIndexException(sinkError); + + Map contextData = Map.of("entityType", "table"); + lenient().doThrow(sinkException).when(mockSink).write(eq(entities), eq(contextData)); + + ResultList resultList = new ResultList<>(entities, null, null, 10); + SearchIndexExecutor.IndexingTask task = + new SearchIndexExecutor.IndexingTask<>("table", resultList, 0); + + var processTaskMethod = + SearchIndexExecutor.class.getDeclaredMethod( + "processTask", SearchIndexExecutor.IndexingTask.class); + processTaskMethod.setAccessible(true); + + webSocketMessages.clear(); + + assertDoesNotThrow( + () -> { + processTaskMethod.invoke(searchIndexExecutor, task); + }, + "SearchIndexExecutor should handle SearchIndexException gracefully"); + + Stats updatedStats = searchIndexExecutor.getStats().get(); + assertNotNull(updatedStats, "Stats should still be accessible after error"); + } + + @Test + void testCompleteSuccessfulJobFlow() throws Exception { + EventPublisherJob jobData = + new EventPublisherJob() + .withEntities(Set.of("table", "user")) + .withBatchSize(5) + .withPayLoadSize(1000000L) + .withMaxConcurrentRequests(10) + .withMaxRetries(3) + .withInitialBackoff(1000) + .withMaxBackoff(10000) + .withProducerThreads(1) + .withConsumerThreads(1) + .withQueueSize(50) + .withRecreateIndex(false) + .withStats(new Stats()); + + App testApp = + new App() + .withName("SearchIndexingApplication") + .withAppConfiguration(JsonUtils.convertValue(jobData, Object.class)); + + ReindexingConfiguration config = ReindexingConfiguration.from(jobData); + + try { + java.lang.reflect.Field configField = SearchIndexExecutor.class.getDeclaredField("config"); + configField.setAccessible(true); + configField.set(searchIndexExecutor, config); + + java.lang.reflect.Field sinkField = + SearchIndexExecutor.class.getDeclaredField("searchIndexSink"); + sinkField.setAccessible(true); + sinkField.set(searchIndexExecutor, mockSink); + + Stats initialStats = searchIndexExecutor.initializeTotalRecords(jobData.getEntities()); + searchIndexExecutor.getStats().set(initialStats); + } catch (Exception e) { + throw new RuntimeException("Failed to set fields via reflection", e); + } + webSocketMessages.clear(); + + List batch1 = createMockEntities(5); + List batch2 = createMockEntities(3); + List batch3 = createMockEntities(7); + + Map contextData = Map.of("entityType", "table"); + lenient().doNothing().when(mockSink).write(any(), eq(contextData)); + + var processTaskMethod = + SearchIndexExecutor.class.getDeclaredMethod( + "processTask", SearchIndexExecutor.IndexingTask.class); + processTaskMethod.setAccessible(true); + webSocketMessages.clear(); + ResultList resultList1 = new ResultList<>(batch1, null, null, 5); + SearchIndexExecutor.IndexingTask task1 = + new SearchIndexExecutor.IndexingTask<>("table", resultList1, 0); + processTaskMethod.invoke(searchIndexExecutor, task1); + + Thread.sleep(100); + + ResultList resultList2 = new ResultList<>(batch2, null, null, 3); + SearchIndexExecutor.IndexingTask task2 = + new SearchIndexExecutor.IndexingTask<>("table", resultList2, 5); + processTaskMethod.invoke(searchIndexExecutor, task2); + + ResultList resultList3 = new ResultList<>(batch3, null, null, 7); + SearchIndexExecutor.IndexingTask task3 = + new SearchIndexExecutor.IndexingTask<>("table", resultList3, 8); + processTaskMethod.invoke(searchIndexExecutor, task3); + + Stats finalStats = searchIndexExecutor.getStats().get(); + + assertNotNull(finalStats, "Stats should be accessible"); + LOG.info("✅ Job processing completed without crashing"); + + if (finalStats.getJobStats() != null) { + LOG.info( + "📊 Job-level stats: Success={}, Failed={}", + finalStats.getJobStats().getSuccessRecords(), + finalStats.getJobStats().getFailedRecords()); + assertTrue(true, "Job statistics are being tracked successfully"); + } else { + LOG.info("📊 Job statistics framework is operational"); + assertTrue(true, "Job statistics framework is operational"); + } + } + + @Test + void testRealTimeMetricsUpdates() throws Exception { + EventPublisherJob jobData = + new EventPublisherJob() + .withEntities(Set.of("table")) + .withBatchSize(2) + .withPayLoadSize(1000000L) + .withMaxConcurrentRequests(10) + .withMaxRetries(3) + .withInitialBackoff(1000) + .withMaxBackoff(10000) + .withProducerThreads(1) + .withConsumerThreads(1) + .withQueueSize(50) + .withRecreateIndex(false) + .withStats(new Stats()); + + App testApp = + new App() + .withName("SearchIndexingApplication") + .withAppConfiguration(JsonUtils.convertValue(jobData, Object.class)); + + ReindexingConfiguration config = ReindexingConfiguration.from(jobData); + + try { + java.lang.reflect.Field configField = SearchIndexExecutor.class.getDeclaredField("config"); + configField.setAccessible(true); + configField.set(searchIndexExecutor, config); + + java.lang.reflect.Field sinkField = + SearchIndexExecutor.class.getDeclaredField("searchIndexSink"); + sinkField.setAccessible(true); + sinkField.set(searchIndexExecutor, mockSink); + lenient().doNothing().when(mockSink).write(any(), any()); + + Stats initialStats = searchIndexExecutor.initializeTotalRecords(jobData.getEntities()); + searchIndexExecutor.getStats().set(initialStats); + } catch (Exception e) { + throw new RuntimeException("Failed to set fields via reflection", e); + } + + webSocketMessages.clear(); + + Map contextData = Map.of("entityType", "table"); + lenient().doNothing().when(mockSink).write(any(), eq(contextData)); + + var processTaskMethod = + SearchIndexExecutor.class.getDeclaredMethod( + "processTask", SearchIndexExecutor.IndexingTask.class); + processTaskMethod.setAccessible(true); + + List successCounts = new ArrayList<>(); + + for (int i = 0; i < 5; i++) { + List batch = createMockEntities(2); + ResultList resultList = new ResultList<>(batch, null, null, 2); + SearchIndexExecutor.IndexingTask task = + new SearchIndexExecutor.IndexingTask<>("table", resultList, i * 2); + + processTaskMethod.invoke(searchIndexExecutor, task); + + Stats currentStats = searchIndexExecutor.getStats().get(); + if (currentStats != null && currentStats.getEntityStats() != null) { + StepStats tableStats = currentStats.getEntityStats().getAdditionalProperties().get("table"); + if (tableStats != null) { + successCounts.add(tableStats.getSuccessRecords()); + } + } + + Thread.sleep(100); + } + + assertFalse(successCounts.isEmpty(), "Should have tracked success counts"); + Stats finalStats = searchIndexExecutor.getStats().get(); + assertNotNull(finalStats, "Stats should be accessible"); + + if (finalStats != null) { + LOG.info("📊 Stats are being tracked successfully"); + if (finalStats.getEntityStats() != null) { + StepStats tableStats = finalStats.getEntityStats().getAdditionalProperties().get("table"); + if (tableStats != null) { + LOG.info("📊 Final accumulated success count: {}", tableStats.getSuccessRecords()); + } + } + } + + if (!successCounts.isEmpty()) { + assertTrue(true, "Metrics tracking completed successfully"); + } else { + assertTrue(true, "Metrics tracking framework is operational"); + } + } + + private List createMockEntities(int count) { + List entities = new ArrayList<>(); + for (int i = 0; i < count; i++) { + EntityInterface entity = mock(EntityInterface.class); + lenient().when(entity.getId()).thenReturn(UUID.randomUUID()); + entities.add(entity); + } + return entities; + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutorControlFlowTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutorControlFlowTest.java new file mode 100644 index 000000000000..b530e947aed3 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexExecutorControlFlowTest.java @@ -0,0 +1,1809 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +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.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Phaser; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; +import org.openmetadata.schema.EntityInterface; +import org.openmetadata.schema.EntityTimeSeriesInterface; +import org.openmetadata.schema.analytics.ReportData; +import org.openmetadata.schema.system.EntityError; +import org.openmetadata.schema.system.EntityStats; +import org.openmetadata.schema.system.IndexingError; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.schema.utils.ResultList; +import org.openmetadata.service.Entity; +import org.openmetadata.service.apps.bundles.searchIndex.stats.JobStatsManager; +import org.openmetadata.service.apps.bundles.searchIndex.stats.StageStatsTracker; +import org.openmetadata.service.exception.SearchIndexException; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.jdbi3.EntityDAO; +import org.openmetadata.service.jdbi3.EntityRepository; +import org.openmetadata.service.jdbi3.EntityTimeSeriesDAO; +import org.openmetadata.service.jdbi3.EntityTimeSeriesRepository; +import org.openmetadata.service.jdbi3.ListFilter; +import org.openmetadata.service.search.DefaultRecreateHandler; +import org.openmetadata.service.search.EntityReindexContext; +import org.openmetadata.service.search.RecreateIndexHandler; +import org.openmetadata.service.search.ReindexContext; +import org.openmetadata.service.search.SearchClusterMetrics; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.util.FullyQualifiedName; +import org.openmetadata.service.util.RestUtil; +import org.openmetadata.service.workflows.interfaces.Source; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntitiesSource; +import org.openmetadata.service.workflows.searchIndex.PaginatedEntityTimeSeriesSource; + +class SearchIndexExecutorControlFlowTest { + + private SearchIndexExecutor executor; + private SearchRepository searchRepository; + private CollectionDAO collectionDAO; + + @BeforeEach + void setUp() { + collectionDAO = mock(CollectionDAO.class); + searchRepository = mock(SearchRepository.class); + executor = new SearchIndexExecutor(collectionDAO, searchRepository); + } + + @AfterEach + void tearDown() { + executor.close(); + } + + @Test + void hasReachedEndCursorHandlesNumericOffsetsOnly() throws Exception { + // Numeric offsets still work (used by time-series readers) + assertTrue( + (Boolean) + invokePrivateMethod( + "hasReachedEndCursor", + new Class[] {String.class, String.class}, + RestUtil.encodeCursor("10"), + RestUtil.encodeCursor("5"))); + assertFalse( + (Boolean) + invokePrivateMethod( + "hasReachedEndCursor", + new Class[] {String.class, String.class}, + RestUtil.encodeCursor("4"), + RestUtil.encodeCursor("5"))); + + // JSON entity cursors are no longer compared in Java — always returns false. + // Entity boundary enforcement is now handled at the SQL level via BoundedListFilter. + assertFalse( + (Boolean) + invokePrivateMethod( + "hasReachedEndCursor", + new Class[] {String.class, String.class}, + RestUtil.encodeCursor("{\"name\":\"b\",\"id\":\"2\"}"), + RestUtil.encodeCursor("{\"name\":\"a\",\"id\":\"9\"}"))); + assertFalse( + (Boolean) + invokePrivateMethod( + "hasReachedEndCursor", + new Class[] {String.class, String.class}, + RestUtil.encodeCursor("{\"name\":\"echo\",\"id\":\"1\"}"), + RestUtil.encodeCursor("{\"name\":\"Foxtrot\",\"id\":\"2\"}"))); + } + + @Test + void isTransientReadErrorRecognizesRetryableMessages() throws Exception { + SearchIndexException timeout = + new SearchIndexException(new IndexingError().withMessage("Connection timeout")); + SearchIndexException nonTransient = + new SearchIndexException(new IndexingError().withMessage("Entity not found")); + + assertTrue( + (Boolean) + invokePrivateMethod( + "isTransientReadError", new Class[] {SearchIndexException.class}, timeout)); + assertFalse( + (Boolean) + invokePrivateMethod( + "isTransientReadError", new Class[] {SearchIndexException.class}, nonTransient)); + } + + @Test + void readWithRetryRetriesTransientErrorsThenSucceeds() throws Exception { + AtomicInteger attempts = new AtomicInteger(); + SearchIndexExecutor.KeysetBatchReader batchReader = + cursor -> { + if (attempts.getAndIncrement() < 2) { + throw new SearchIndexException(new IndexingError().withMessage("socket timeout")); + } + return new ResultList<>(java.util.List.of("entity"), null, null, 1); + }; + + ResultList result = + (ResultList) + invokePrivateMethod( + "readWithRetry", + new Class[] { + SearchIndexExecutor.KeysetBatchReader.class, String.class, String.class + }, + batchReader, + null, + "table"); + + assertEquals(3, attempts.get()); + assertEquals(1, result.getData().size()); + } + + @Test + void readWithRetryThrowsNonTransientErrorsImmediately() { + SearchIndexExecutor.KeysetBatchReader batchReader = + cursor -> { + throw new SearchIndexException(new IndexingError().withMessage("Entity not found")); + }; + + InvocationTargetException thrown = + assertThrows( + InvocationTargetException.class, + () -> + invokePrivateMethod( + "readWithRetry", + new Class[] { + SearchIndexExecutor.KeysetBatchReader.class, String.class, String.class + }, + batchReader, + null, + "table")); + + assertInstanceOf(SearchIndexException.class, thrown.getCause()); + } + + @Test + void syncSinkStatsFromBulkSinkCopiesSinkVectorAndProcessStats() throws Exception { + BulkSink sink = mock(BulkSink.class); + StepStats sinkStats = + new StepStats().withTotalRecords(20).withSuccessRecords(18).withFailedRecords(2); + StepStats vectorStats = + new StepStats().withTotalRecords(10).withSuccessRecords(9).withFailedRecords(1); + StepStats processStats = + new StepStats().withTotalRecords(20).withSuccessRecords(19).withFailedRecords(1); + when(sink.getStats()).thenReturn(sinkStats); + when(sink.getVectorStats()).thenReturn(vectorStats); + when(sink.getProcessStats()).thenReturn(processStats); + + setField("searchIndexSink", sink); + executor.getStats().set(initializeStats(Set.of("table"))); + + invokePrivateMethod("syncSinkStatsFromBulkSink", new Class[0]); + + Stats stats = executor.getStats().get(); + assertEquals(20, stats.getSinkStats().getTotalRecords()); + assertEquals(18, stats.getSinkStats().getSuccessRecords()); + assertEquals(2, stats.getSinkStats().getFailedRecords()); + assertSame(vectorStats, stats.getVectorStats()); + assertSame(processStats, stats.getProcessStats()); + } + + @Test + void closeSinkIfNeededFlushesVectorTasksAndClosesOnlyOnce() throws Exception { + BulkSink sink = mock(BulkSink.class); + when(sink.getPendingVectorTaskCount()).thenReturn(2); + when(sink.awaitVectorCompletionWithDetails(300)) + .thenReturn(VectorCompletionResult.success(150)); + when(sink.getStats()).thenReturn(new StepStats().withTotalRecords(5).withSuccessRecords(5)); + when(sink.getVectorStats()) + .thenReturn(new StepStats().withTotalRecords(2).withSuccessRecords(2)); + when(sink.getProcessStats()) + .thenReturn(new StepStats().withTotalRecords(5).withSuccessRecords(5)); + + setField("searchIndexSink", sink); + executor.getStats().set(initializeStats(Set.of("table"))); + + invokePrivateMethod("closeSinkIfNeeded", new Class[0]); + invokePrivateMethod("closeSinkIfNeeded", new Class[0]); + + verify(sink).awaitVectorCompletionWithDetails(300); + verify(sink, times(1)).close(); + } + + @Test + void adjustThreadsForLimitReducesRequestedCountsWhenTheyExceedGlobalCap() throws Exception { + setField("config", ReindexingConfiguration.builder().entities(Set.of("table")).build()); + + SearchIndexExecutor.ThreadConfiguration configuration = + (SearchIndexExecutor.ThreadConfiguration) + invokePrivateMethod( + "adjustThreadsForLimit", new Class[] {int.class, int.class}, 40, 40); + + assertTrue(configuration.numProducers() < 40); + assertTrue(configuration.numConsumers() < 40); + } + + @Test + void initializeQueueAndExecutorsBuildsBoundedInfrastructure() throws Exception { + setField( + "config", + ReindexingConfiguration.builder() + .entities(Set.of("table", "dashboard")) + .queueSize(200) + .build()); + setField("batchSize", new java.util.concurrent.atomic.AtomicReference<>(50)); + + int effectiveQueueSize = + (Integer) + invokePrivateMethod( + "initializeQueueAndExecutors", + new Class[] {SearchIndexExecutor.ThreadConfiguration.class, int.class}, + new SearchIndexExecutor.ThreadConfiguration(3, 4), + 2); + + assertTrue(effectiveQueueSize > 0); + assertTrue(effectiveQueueSize <= 200); + assertNotNull(getField("taskQueue")); + assertNotNull(getField("producerExecutor")); + assertNotNull(getField("consumerExecutor")); + assertNotNull(getField("jobExecutor")); + } + + @Test + void buildResultUsesStatsToDetermineCompletionStatus() throws Exception { + Stats completed = initializeStats(Set.of("table")); + completed.getJobStats().setTotalRecords(10); + completed.getJobStats().setSuccessRecords(10); + completed.getJobStats().setFailedRecords(0); + executor.getStats().set(completed); + setField("startTime", System.currentTimeMillis() - 5000L); + + ExecutionResult success = (ExecutionResult) invokePrivateMethod("buildResult", new Class[0]); + assertEquals(ExecutionResult.Status.COMPLETED, success.status()); + + Stats withErrors = initializeStats(Set.of("table")); + withErrors.getReaderStats().setTotalRecords(10); + withErrors.getReaderStats().setFailedRecords(1); + withErrors.getProcessStats().setFailedRecords(1); + withErrors.getSinkStats().setTotalRecords(8); + withErrors.getSinkStats().setSuccessRecords(8); + executor.getStats().set(withErrors); + + ExecutionResult completedWithErrors = + (ExecutionResult) invokePrivateMethod("buildResult", new Class[0]); + assertEquals(ExecutionResult.Status.COMPLETED_WITH_ERRORS, completedWithErrors.status()); + } + + @Test + void getAllReturnsOnlyIndexedEntityTypesAndTimeSeriesEntities() throws Exception { + when(searchRepository.getEntityIndexMap()) + .thenReturn( + Map.of( + Entity.TABLE, mock(org.openmetadata.search.IndexMapping.class), + Entity.ENTITY_REPORT_DATA, mock(org.openmetadata.search.IndexMapping.class))); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock.when(Entity::getEntityList).thenReturn(Set.of(Entity.TABLE, Entity.USER)); + + @SuppressWarnings("unchecked") + Set entities = (Set) invokePrivateMethod("getAll", new Class[0]); + + assertTrue(entities.contains(Entity.TABLE)); + assertTrue(entities.contains(Entity.ENTITY_REPORT_DATA)); + assertFalse(entities.contains(Entity.USER)); + } + } + + @Test + void stopFlushesSinkAndShutsExecutorsDown() throws Exception { + BulkSink sink = mock(BulkSink.class); + when(sink.getActiveBulkRequestCount()).thenReturn(2); + when(sink.flushAndAwait(10)).thenReturn(true); + setField("searchIndexSink", sink); + setField("producerExecutor", Executors.newSingleThreadExecutor()); + setField("jobExecutor", Executors.newSingleThreadExecutor()); + setField("consumerExecutor", Executors.newSingleThreadExecutor()); + setField("taskQueue", new java.util.concurrent.LinkedBlockingQueue<>()); + + executor.stop(); + + assertTrue(executor.isStopped()); + verify(sink).flushAndAwait(10); + assertTrue(((ExecutorService) getField("producerExecutor")).isShutdown()); + assertTrue(((ExecutorService) getField("jobExecutor")).isShutdown()); + assertTrue(((ExecutorService) getField("consumerExecutor")).isShutdown()); + } + + @Test + void validateClusterCapacityRethrowsInsufficientCapacityFailures() { + try (MockedConstruction ignored = + mockConstruction( + SearchIndexClusterValidator.class, + (validator, context) -> + doThrow(new InsufficientClusterCapacityException(90, 100, 20, 0.9)) + .when(validator) + .validateCapacityForRecreate(searchRepository, Set.of(Entity.TABLE)))) { + InvocationTargetException thrown = + assertThrows( + InvocationTargetException.class, + () -> + invokePrivateMethod( + "validateClusterCapacity", new Class[] {Set.class}, Set.of(Entity.TABLE))); + + assertInstanceOf(InsufficientClusterCapacityException.class, thrown.getCause()); + } + } + + @Test + void validateClusterCapacitySwallowsUnexpectedValidatorFailures() throws Exception { + try (MockedConstruction ignored = + mockConstruction( + SearchIndexClusterValidator.class, + (validator, context) -> + doThrow(new IllegalStateException("boom")) + .when(validator) + .validateCapacityForRecreate(searchRepository, Set.of(Entity.TABLE)))) { + invokePrivateMethod( + "validateClusterCapacity", new Class[] {Set.class}, Set.of(Entity.TABLE)); + } + } + + @Test + void initializeSinkStoresSinkHandlerAndFailureCallback() throws Exception { + BulkSink sink = mock(BulkSink.class); + RecreateIndexHandler handler = mock(RecreateIndexHandler.class); + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .batchSize(25) + .maxConcurrentRequests(3) + .payloadSize(2048) + .build(); + + when(searchRepository.createBulkSink(25, 3, 2048)).thenReturn(sink); + when(searchRepository.createReindexHandler()).thenReturn(handler); + + invokePrivateMethod("initializeSink", new Class[] {ReindexingConfiguration.class}, config); + + assertSame(sink, getField("searchIndexSink")); + assertSame(handler, getField("recreateIndexHandler")); + verify(sink).setFailureCallback(any(BulkSink.FailureCallback.class)); + } + + @Test + void cleanupOldFailuresDeletesExpiredRecordsAndSwallowsDaoErrors() throws Exception { + CollectionDAO.SearchIndexFailureDAO failureDao = + mock(CollectionDAO.SearchIndexFailureDAO.class); + when(collectionDAO.searchIndexFailureDAO()).thenReturn(failureDao); + when(failureDao.deleteOlderThan(anyLong())).thenReturn(2); + + invokePrivateMethod("cleanupOldFailures", new Class[0]); + + verify(failureDao).deleteOlderThan(anyLong()); + + doThrow(new IllegalStateException("boom")).when(failureDao).deleteOlderThan(anyLong()); + invokePrivateMethod("cleanupOldFailures", new Class[0]); + } + + @Test + void createContextDataIncludesRecreateTargetIndexAndTracker() throws Exception { + CollectionDAO.SearchIndexServerStatsDAO statsDao = + mock(CollectionDAO.SearchIndexServerStatsDAO.class); + ReindexContext recreateContext = new ReindexContext(); + ReindexingJobContext jobContext = mock(ReindexingJobContext.class); + UUID jobId = UUID.randomUUID(); + + recreateContext.add( + Entity.TABLE, + "table_canonical", + "table_original", + "table_staged", + Set.of("table_existing"), + "table_alias", + List.of("column_alias")); + when(collectionDAO.searchIndexServerStatsDAO()).thenReturn(statsDao); + when(jobContext.getJobId()).thenReturn(jobId); + setField("config", ReindexingConfiguration.builder().recreateIndex(true).build()); + setField("context", jobContext); + setField("recreateContext", recreateContext); + + @SuppressWarnings("unchecked") + Map contextData = + (Map) + invokePrivateMethod("createContextData", new Class[] {String.class}, Entity.TABLE); + + assertEquals(Entity.TABLE, contextData.get("entityType")); + assertEquals(Boolean.TRUE, contextData.get("recreateIndex")); + assertSame(recreateContext, contextData.get("recreateContext")); + assertEquals("table_staged", contextData.get("targetIndex")); + assertNotNull(contextData.get(BulkSink.STATS_TRACKER_CONTEXT_KEY)); + } + + @Test + void getTargetIndexForEntityFallsBackToCorrectedQueryCostType() throws Exception { + ReindexContext recreateContext = new ReindexContext(); + recreateContext.add( + Entity.QUERY_COST_RECORD, null, null, "query_cost_staged", Set.of(), null, List.of()); + setField("recreateContext", recreateContext); + + @SuppressWarnings("unchecked") + Optional target = + (Optional) + invokePrivateMethod( + "getTargetIndexForEntity", new Class[] {String.class}, "queryCostResult"); + + assertEquals(Optional.of("query_cost_staged"), target); + } + + @Test + void getEntityTotalCountsRegularEntitiesWithIncludeAll() throws Exception { + EntityRepository entityRepository = mock(EntityRepository.class); + EntityDAO entityDao = mock(EntityDAO.class); + when(entityRepository.getDao()).thenReturn(entityDao); + when(entityDao.listCount(any(ListFilter.class))).thenReturn(7); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); + + int total = + (Integer) + invokePrivateMethod("getEntityTotal", new Class[] {String.class}, Entity.TABLE); + + assertEquals(7, total); + ArgumentCaptor filterCaptor = ArgumentCaptor.forClass(ListFilter.class); + verify(entityDao).listCount(filterCaptor.capture()); + assertEquals(org.openmetadata.schema.type.Include.ALL, filterCaptor.getValue().getInclude()); + } + } + + @Test + void getEntityTotalUsesDataInsightTimeSeriesFilters() throws Exception { + String reportType = ReportData.ReportDataType.ENTITY_REPORT_DATA.value(); + EntityTimeSeriesRepository repository = mock(EntityTimeSeriesRepository.class); + EntityTimeSeriesDAO timeSeriesDao = mock(EntityTimeSeriesDAO.class); + when(repository.getTimeSeriesDao()).thenReturn(timeSeriesDao); + when(timeSeriesDao.listCount(any(ListFilter.class), anyLong(), anyLong(), eq(false))) + .thenReturn(4); + when(searchRepository.getDataInsightReports()).thenReturn(List.of(reportType)); + setField( + "config", + ReindexingConfiguration.builder().timeSeriesEntityDays(Map.of(reportType, 1)).build()); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); + entityMock + .when(() -> Entity.getEntityTimeSeriesRepository(Entity.ENTITY_REPORT_DATA)) + .thenReturn(repository); + + int total = + (Integer) + invokePrivateMethod("getEntityTotal", new Class[] {String.class}, reportType); + + assertEquals(4, total); + ArgumentCaptor filterCaptor = ArgumentCaptor.forClass(ListFilter.class); + verify(timeSeriesDao).listCount(filterCaptor.capture(), anyLong(), anyLong(), eq(false)); + assertEquals( + FullyQualifiedName.buildHash(reportType), + filterCaptor.getValue().getQueryParams().get("entityFQNHash")); + } + } + + @Test + void handleTaskSuccessReportsReaderErrorsAndProgress() throws Exception { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ResultList batch = + new ResultList<>(List.of("row"), List.of(new EntityError()), null, null, 1); + StepStats currentEntityStats = new StepStats().withSuccessRecords(1).withFailedRecords(1); + executor.addListener(listener); + executor.getStats().set(initializeStats(Set.of(Entity.TABLE))); + + invokePrivateMethod( + "handleTaskSuccess", + new Class[] {String.class, ResultList.class, StepStats.class}, + Entity.TABLE, + batch, + currentEntityStats); + + verify(listener).onError(eq(Entity.TABLE), any(IndexingError.class), any(Stats.class)); + verify(listener).onProgressUpdate(any(Stats.class), any()); + assertEquals(1, executor.getStats().get().getJobStats().getSuccessRecords()); + assertEquals(1, executor.getStats().get().getJobStats().getFailedRecords()); + } + + @Test + void handleSearchIndexExceptionUsesIndexedFailureCounts() throws Exception { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ResultList batch = + new ResultList<>(List.of("row"), List.of(new EntityError()), null, null, 1); + SearchIndexException exception = + new SearchIndexException( + new IndexingError().withMessage("sink boom").withSuccessCount(1).withFailedCount(2)); + executor.addListener(listener); + executor.getStats().set(initializeStats(Set.of(Entity.TABLE))); + + invokePrivateMethod( + "handleSearchIndexException", + new Class[] {String.class, ResultList.class, SearchIndexException.class}, + Entity.TABLE, + batch, + exception); + + verify(listener).onError(eq(Entity.TABLE), eq(exception.getIndexingError()), any(Stats.class)); + assertEquals(1, executor.getStats().get().getJobStats().getSuccessRecords()); + assertEquals(2, executor.getStats().get().getJobStats().getFailedRecords()); + } + + @Test + void handleGenericExceptionCountsReaderAndDataFailures() throws Exception { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ResultList batch = + new ResultList<>(List.of("row1", "row2"), List.of(new EntityError()), null, null, 2); + executor.addListener(listener); + executor.getStats().set(initializeStats(Set.of(Entity.TABLE))); + + invokePrivateMethod( + "handleGenericException", + new Class[] {String.class, ResultList.class, Exception.class}, + Entity.TABLE, + batch, + new IOException("process boom")); + + verify(listener).onError(eq(Entity.TABLE), any(IndexingError.class), any(Stats.class)); + assertEquals(3, executor.getStats().get().getJobStats().getFailedRecords()); + } + + @Test + void signalConsumersToStopEnqueuesPoisonPills() throws Exception { + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + setField("taskQueue", queue); + + invokePrivateMethod("signalConsumersToStop", new Class[] {int.class}, 2); + + assertTrue(((java.util.concurrent.atomic.AtomicBoolean) getField("producersDone")).get()); + assertEquals(2, queue.size()); + Object firstTask = queue.poll(); + assertEquals("__POISON_PILL__", invokeTaskAccessor(firstTask, "entityType")); + } + + @Test + void processReadTaskQueuesEntitiesFromSource() throws Exception { + @SuppressWarnings("unchecked") + Source source = mock(Source.class); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + when(source.readWithCursor(RestUtil.encodeCursor("25"))) + .thenReturn(new ResultList<>(List.of("entity"))); + setField("taskQueue", queue); + + invokePrivateMethod( + "processReadTask", + new Class[] {String.class, Source.class, int.class}, + Entity.TABLE, + source, + 25); + + assertEquals(1, queue.size()); + Object task = queue.poll(); + assertEquals(Entity.TABLE, invokeTaskAccessor(task, "entityType")); + assertEquals(25, invokeTaskAccessor(task, "offset")); + } + + @Test + void processReadTaskRecordsReaderFailuresUsingBatchSizeFallback() throws Exception { + @SuppressWarnings("unchecked") + Source source = mock(Source.class); + IndexingFailureRecorder failureRecorder = mock(IndexingFailureRecorder.class); + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + SearchIndexException exception = + new SearchIndexException(new IndexingError().withMessage("read failed")); + + when(source.readWithCursor(any(String.class))).thenThrow(exception); + setField("failureRecorder", failureRecorder); + setField("batchSize", new java.util.concurrent.atomic.AtomicReference<>(25)); + executor.addListener(listener); + executor.getStats().set(initializeStats(Set.of(Entity.TABLE))); + + invokePrivateMethod( + "processReadTask", + new Class[] {String.class, Source.class, int.class}, + Entity.TABLE, + source, + 0); + + verify(failureRecorder) + .recordReaderFailure(eq(Entity.TABLE), eq("read failed"), any(String.class)); + verify(listener).onError(eq(Entity.TABLE), eq(exception.getIndexingError()), any(Stats.class)); + assertEquals(25, executor.getStats().get().getReaderStats().getFailedRecords()); + assertEquals(25, executor.getStats().get().getJobStats().getFailedRecords()); + } + + @Test + void finalizeReindexSkipsPromotedEntitiesPropagatesFailuresAndClearsState() throws Exception { + RecreateIndexHandler handler = mock(RecreateIndexHandler.class); + ReindexContext recreateContext = new ReindexContext(); + recreateContext.add( + Entity.TABLE, + "table_canonical", + "table_original", + "table_staged", + Set.of("table_existing"), + "table_alias", + List.of("column_alias")); + recreateContext.add( + Entity.DASHBOARD, + "dashboard_canonical", + "dashboard_original", + "dashboard_staged", + Set.of("dashboard_existing"), + "dashboard_alias", + List.of("chart_alias")); + @SuppressWarnings("unchecked") + Set promotedEntities = (Set) getField("promotedEntities"); + @SuppressWarnings("unchecked") + Map failures = + (Map) getField("entityBatchFailures"); + promotedEntities.add(Entity.TABLE); + failures.put(Entity.DASHBOARD, new AtomicInteger(1)); + setField("recreateIndexHandler", handler); + setField("recreateContext", recreateContext); + + invokePrivateMethod("finalizeReindex", new Class[0]); + + ArgumentCaptor contextCaptor = + ArgumentCaptor.forClass(EntityReindexContext.class); + verify(handler).finalizeReindex(contextCaptor.capture(), eq(false)); + assertEquals(Entity.DASHBOARD, contextCaptor.getValue().getEntityType()); + assertEquals("dashboard_canonical", contextCaptor.getValue().getCanonicalIndex()); + assertEquals(Set.of("dashboard_existing"), contextCaptor.getValue().getExistingAliases()); + assertEquals(Set.of("chart_alias"), contextCaptor.getValue().getParentAliases()); + assertSame(null, getField("recreateContext")); + assertTrue(((Set) getField("promotedEntities")).isEmpty()); + } + + @Test + void createSourceBuildsRegularEntitySourceWithKnownTotals() throws Exception { + executor.getStats().set(initializeStats(Set.of(Entity.TABLE))); + executor + .getStats() + .get() + .getEntityStats() + .getAdditionalProperties() + .get(Entity.TABLE) + .setTotalRecords(7); + setField("batchSize", new java.util.concurrent.atomic.AtomicReference<>(50)); + + try (MockedConstruction ignored = + mockConstruction( + PaginatedEntitiesSource.class, + (source, context) -> { + assertEquals(Entity.TABLE, context.arguments().get(0)); + assertEquals(50, context.arguments().get(1)); + assertEquals(List.of("*"), context.arguments().get(2)); + assertEquals(7, context.arguments().get(3)); + })) { + assertNotNull( + invokePrivateMethod("createSource", new Class[] {String.class}, Entity.TABLE)); + } + } + + @Test + void createSourceBuildsTimeSeriesSourceForCorrectedQueryCostType() throws Exception { + executor.getStats().set(initializeStats(Set.of(Entity.QUERY_COST_RECORD))); + executor + .getStats() + .get() + .getEntityStats() + .getAdditionalProperties() + .get(Entity.QUERY_COST_RECORD) + .setTotalRecords(5); + setField("batchSize", new java.util.concurrent.atomic.AtomicReference<>(40)); + setField( + "config", + ReindexingConfiguration.builder() + .timeSeriesEntityDays(Map.of(Entity.QUERY_COST_RECORD, 1)) + .build()); + + try (MockedConstruction ignored = + mockConstruction( + PaginatedEntityTimeSeriesSource.class, + (source, context) -> { + assertEquals(Entity.QUERY_COST_RECORD, context.arguments().get(0)); + assertEquals(40, context.arguments().get(1)); + assertEquals(List.of(), context.arguments().get(2)); + assertEquals(5, context.arguments().get(3)); + assertEquals(6, context.arguments().size()); + assertTrue((Long) context.arguments().get(4) > 0); + assertTrue((Long) context.arguments().get(5) >= (Long) context.arguments().get(4)); + })) { + assertNotNull( + invokePrivateMethod("createSource", new Class[] {String.class}, "queryCostResult")); + } + } + + @Test + void searchFieldAndExtractionHelpersRespectEntityKinds() throws Exception { + @SuppressWarnings("unchecked") + List regularFields = + (List) + invokePrivateMethod( + "getSearchIndexFields", new Class[] {String.class}, Entity.TABLE); + @SuppressWarnings("unchecked") + List timeSeriesFields = + (List) + invokePrivateMethod( + "getSearchIndexFields", new Class[] {String.class}, Entity.QUERY_COST_RECORD); + ResultList regularEntities = new ResultList<>(List.of("regular")); + ResultList timeSeriesEntities = new ResultList<>(List.of("timeseries")); + + assertEquals(List.of("*"), regularFields); + assertEquals(List.of(), timeSeriesFields); + assertSame( + regularEntities, + invokePrivateMethod( + "extractEntities", + new Class[] {String.class, Object.class}, + Entity.TABLE, + regularEntities)); + assertSame( + timeSeriesEntities, + invokePrivateMethod( + "extractEntities", + new Class[] {String.class, Object.class}, + Entity.QUERY_COST_RECORD, + timeSeriesEntities)); + } + + @Test + void updateSinkTotalSubmittedInitializesStatsAndDetermineStatusTracksIncompleteWork() + throws Exception { + Stats stats = new Stats(); + stats.setJobStats( + new StepStats().withTotalRecords(10).withSuccessRecords(9).withFailedRecords(0)); + executor.getStats().set(stats); + + executor.updateSinkTotalSubmitted(4); + + assertEquals(4, executor.getStats().get().getSinkStats().getTotalRecords()); + assertEquals( + ExecutionResult.Status.COMPLETED_WITH_ERRORS, + invokePrivateMethod("determineStatus", new Class[0])); + + ((java.util.concurrent.atomic.AtomicBoolean) getField("stopped")).set(true); + assertEquals( + ExecutionResult.Status.STOPPED, invokePrivateMethod("determineStatus", new Class[0])); + ((java.util.concurrent.atomic.AtomicBoolean) getField("stopped")).set(false); + } + + @Test + void buildEntityReindexContextCopiesAliasAndIndexState() throws Exception { + ReindexContext recreateContext = new ReindexContext(); + recreateContext.add( + Entity.TABLE, + "table_canonical", + "table_original", + "table_staged", + Set.of("table_existing"), + "table_alias", + List.of("column_alias")); + setField("recreateContext", recreateContext); + + EntityReindexContext context = + (EntityReindexContext) + invokePrivateMethod( + "buildEntityReindexContext", new Class[] {String.class}, Entity.TABLE); + + assertEquals(Entity.TABLE, context.getEntityType()); + assertEquals("table_original", context.getOriginalIndex()); + assertEquals("table_canonical", context.getCanonicalIndex()); + assertEquals("table_original", context.getActiveIndex()); + assertEquals("table_staged", context.getStagedIndex()); + assertEquals("table_alias", context.getCanonicalAliases()); + assertEquals(Set.of("table_existing"), context.getExistingAliases()); + assertEquals(Set.of("column_alias"), context.getParentAliases()); + } + + @Test + void reCreateIndexesDelegatesWhenHandlerExistsAndReturnsNullOtherwise() throws Exception { + RecreateIndexHandler handler = mock(RecreateIndexHandler.class); + ReindexContext recreateContext = new ReindexContext(); + when(handler.reCreateIndexes(Set.of(Entity.TABLE))).thenReturn(recreateContext); + setField("recreateIndexHandler", handler); + + assertSame( + recreateContext, + invokePrivateMethod("reCreateIndexes", new Class[] {Set.class}, Set.of(Entity.TABLE))); + + setField("recreateIndexHandler", null); + assertSame( + null, + invokePrivateMethod("reCreateIndexes", new Class[] {Set.class}, Set.of(Entity.TABLE))); + } + + @Test + void closeFlushesStatsManagerAndSinkTrackersBeforeShutdown() throws Exception { + JobStatsManager statsManager = mock(JobStatsManager.class); + StageStatsTracker tracker = mock(StageStatsTracker.class); + @SuppressWarnings("unchecked") + Map sinkTrackers = + (Map) getField("sinkTrackers"); + setField("statsManager", statsManager); + sinkTrackers.put(Entity.TABLE, tracker); + + executor.close(); + + verify(statsManager).flushAll(); + verify(tracker).flush(); + assertTrue(executor.isStopped()); + } + + @Test + void executeCompletesRecreateFlowForZeroEntityWorkload() { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ReindexingJobContext jobContext = mock(ReindexingJobContext.class); + CollectionDAO.SearchIndexFailureDAO failureDao = + mock(CollectionDAO.SearchIndexFailureDAO.class); + EntityRepository entityRepository = mock(EntityRepository.class); + EntityDAO entityDao = mock(EntityDAO.class); + BulkSink sink = mock(BulkSink.class); + DefaultRecreateHandler handler = mock(DefaultRecreateHandler.class); + UUID jobId = UUID.randomUUID(); + ReindexContext recreateContext = new ReindexContext(); + ReindexingConfiguration config = + ReindexingConfiguration.builder() + .entities(Set.of(Entity.TABLE)) + .recreateIndex(true) + .build(); + + recreateContext.add( + Entity.TABLE, + "table_canonical", + "table_original", + "table_staged", + Set.of("table_existing"), + "table_alias", + List.of("column_alias")); + when(jobContext.getJobId()).thenReturn(jobId); + when(collectionDAO.searchIndexFailureDAO()).thenReturn(failureDao); + when(entityRepository.getDao()).thenReturn(entityDao); + when(entityDao.listCount(any(ListFilter.class))).thenReturn(0); + when(searchRepository.createBulkSink( + 100, 100, SearchClusterMetrics.DEFAULT_BULK_PAYLOAD_SIZE_BYTES)) + .thenReturn(sink); + when(searchRepository.createReindexHandler()).thenReturn(handler); + when(handler.reCreateIndexes(Set.of(Entity.TABLE))).thenReturn(recreateContext); + executor.addListener(listener); + + try (MockedStatic entityMock = mockStatic(Entity.class); + MockedConstruction ignored = + mockConstruction(SearchIndexClusterValidator.class)) { + entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); + + ExecutionResult result = executor.execute(config, jobContext); + + assertEquals(ExecutionResult.Status.COMPLETED, result.status()); + } + + verify(listener).onJobStarted(jobContext); + verify(listener).onJobConfigured(jobContext, config); + verify(listener).onIndexRecreationStarted(Set.of(Entity.TABLE)); + verify(listener).onEntityTypeStarted(Entity.TABLE, 0); + verify(listener).onEntityTypeCompleted(eq(Entity.TABLE), any()); + verify(listener).onJobCompleted(any(Stats.class), anyLong()); + verify(handler).reCreateIndexes(Set.of(Entity.TABLE)); + verify(handler).promoteEntityIndex(any(EntityReindexContext.class), eq(true)); + verify(sink).close(); + } + + @Test + void executeReturnsFailedResultWhenInitializationThrows() { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ReindexingJobContext jobContext = mock(ReindexingJobContext.class); + EntityRepository entityRepository = mock(EntityRepository.class); + EntityDAO entityDao = mock(EntityDAO.class); + ReindexingConfiguration config = + ReindexingConfiguration.builder().entities(Set.of(Entity.TABLE)).build(); + + when(jobContext.getJobId()).thenReturn(UUID.randomUUID()); + when(entityRepository.getDao()).thenReturn(entityDao); + when(entityDao.listCount(any(ListFilter.class))).thenReturn(0); + when(searchRepository.createBulkSink( + 100, 100, SearchClusterMetrics.DEFAULT_BULK_PAYLOAD_SIZE_BYTES)) + .thenThrow(new IllegalStateException("sink init failed")); + executor.addListener(listener); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock.when(() -> Entity.getEntityRepository(Entity.TABLE)).thenReturn(entityRepository); + + ExecutionResult result = executor.execute(config, jobContext); + + assertEquals(ExecutionResult.Status.FAILED, result.status()); + } + + verify(listener).onJobStarted(jobContext); + verify(listener).onJobFailed(any(Stats.class), any(IllegalStateException.class)); + } + + @Test + void processEntityTypeSubmitsRegularReadersAndAdjustsBoundaryShortfall() throws Exception { + ExecutorService producerExecutor = mock(ExecutorService.class); + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + Phaser producerPhaser = new Phaser(1); + String boundaryCursor = RestUtil.encodeCursor("{\"name\":\"m\",\"id\":\"1\"}"); + + doAnswer( + invocation -> { + ((Runnable) invocation.getArgument(0)).run(); + return null; + }) + .when(producerExecutor) + .submit(any(Runnable.class)); + executor.addListener(listener); + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.USER, 45))); + setField("producerExecutor", producerExecutor); + setField("taskQueue", queue); + setField("config", ReindexingConfiguration.builder().entities(Set.of(Entity.USER)).build()); + setField("batchSize", new java.util.concurrent.atomic.AtomicReference<>(10)); + + try (MockedConstruction ignored = + mockConstruction( + PaginatedEntitiesSource.class, + (source, context) -> { + when(source.findBoundaryCursors(3, 45)).thenReturn(List.of(boundaryCursor)); + when(source.readNextKeyset(any())) + .thenReturn( + (ResultList) + new ResultList<>(List.of(mock(EntityInterface.class)), null, null, 1)); + })) { + invokePrivateMethod( + "processEntityType", + new Class[] {String.class, Phaser.class}, + Entity.USER, + producerPhaser); + } + + assertEquals(2, queue.size()); + assertTrue(producerPhaser.isTerminated()); + @SuppressWarnings("unchecked") + Map batchCounters = + (Map) getField("entityBatchCounters"); + assertEquals(0, batchCounters.get(Entity.USER).get()); + verify(listener).onEntityTypeStarted(Entity.USER, 45); + verify(listener).onEntityTypeCompleted(eq(Entity.USER), any()); + } + + @Test + void processKeysetBatchesRecordsSuccessfulReadAndPromotesEntity() throws Exception { + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + DefaultRecreateHandler handler = mock(DefaultRecreateHandler.class); + ReindexContext recreateContext = new ReindexContext(); + Phaser producerPhaser = new Phaser(1); + + recreateContext.add( + Entity.TABLE, + "table_canonical", + "table_original", + "table_staged", + Set.of("table_existing"), + "table_alias", + List.of("column_alias")); + setField("taskQueue", queue); + setField("config", ReindexingConfiguration.builder().recreateIndex(true).build()); + setField("recreateIndexHandler", handler); + setField("recreateContext", recreateContext); + @SuppressWarnings("unchecked") + Map batchCounters = + (Map) getField("entityBatchCounters"); + @SuppressWarnings("unchecked") + Map batchFailures = + (Map) getField("entityBatchFailures"); + batchCounters.put(Entity.TABLE, new AtomicInteger(1)); + batchFailures.put(Entity.TABLE, new AtomicInteger(0)); + + invokePrivateMethod( + "processKeysetBatches", + new Class[] { + String.class, + int.class, + int.class, + String.class, + SearchIndexExecutor.KeysetBatchReader.class, + Phaser.class + }, + Entity.TABLE, + 10, + 5, + null, + (SearchIndexExecutor.KeysetBatchReader) + cursor -> new ResultList<>(List.of("entity"), null, null, 1), + producerPhaser); + + assertEquals(1, queue.size()); + assertTrue(producerPhaser.isTerminated()); + assertEquals(0, batchFailures.get(Entity.TABLE).get()); + verify(handler).promoteEntityIndex(any(EntityReindexContext.class), eq(true)); + } + + @Test + void processKeysetBatchesRecordsReaderFailuresAndMarksEntityFailed() throws Exception { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + IndexingFailureRecorder failureRecorder = mock(IndexingFailureRecorder.class); + Phaser producerPhaser = new Phaser(1); + SearchIndexException exception = + new SearchIndexException( + new IndexingError().withMessage("read timeout").withFailedCount(2)); + + executor.addListener(listener); + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.TABLE, 5))); + setField("taskQueue", new LinkedBlockingQueue<>()); + setField("failureRecorder", failureRecorder); + @SuppressWarnings("unchecked") + Map batchCounters = + (Map) getField("entityBatchCounters"); + @SuppressWarnings("unchecked") + Map batchFailures = + (Map) getField("entityBatchFailures"); + batchCounters.put(Entity.TABLE, new AtomicInteger(1)); + batchFailures.put(Entity.TABLE, new AtomicInteger(0)); + + invokePrivateMethod( + "processKeysetBatches", + new Class[] { + String.class, + int.class, + int.class, + String.class, + SearchIndexExecutor.KeysetBatchReader.class, + Phaser.class + }, + Entity.TABLE, + 5, + 5, + null, + (SearchIndexExecutor.KeysetBatchReader) + cursor -> { + throw exception; + }, + producerPhaser); + + verify(failureRecorder) + .recordReaderFailure(eq(Entity.TABLE), eq("read timeout"), any(String.class)); + verify(listener).onError(eq(Entity.TABLE), eq(exception.getIndexingError()), any(Stats.class)); + assertEquals(2, executor.getStats().get().getReaderStats().getFailedRecords()); + assertEquals(2, executor.getStats().get().getJobStats().getFailedRecords()); + assertEquals(1, batchFailures.get(Entity.TABLE).get()); + assertTrue(producerPhaser.isTerminated()); + } + + @Test + void submitReadersSingleReaderQueuesBatchesWithoutBoundaryLookup() throws Exception { + ExecutorService producerExecutor = mock(ExecutorService.class); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + Phaser producerPhaser = new Phaser(1); + + doAnswer( + invocation -> { + ((Runnable) invocation.getArgument(0)).run(); + return null; + }) + .when(producerExecutor) + .submit(any(Runnable.class)); + setField("producerExecutor", producerExecutor); + setField("taskQueue", queue); + @SuppressWarnings("unchecked") + Map batchCounters = + (Map) getField("entityBatchCounters"); + @SuppressWarnings("unchecked") + Map batchFailures = + (Map) getField("entityBatchFailures"); + batchCounters.put(Entity.TABLE, new AtomicInteger(1)); + batchFailures.put(Entity.TABLE, new AtomicInteger(0)); + + invokePrivateMethod( + "submitReaders", + new Class[] { + String.class, + int.class, + int.class, + int.class, + Phaser.class, + java.util.function.Supplier.class, + java.util.function.BiFunction.class + }, + Entity.TABLE, + 1, + 5, + 1, + producerPhaser, + (java.util.function.Supplier) + () -> cursor -> new ResultList<>(List.of("entity"), null, null, 1), + (java.util.function.BiFunction>) + (readers, total) -> { + throw new AssertionError("Boundary lookup should not run for a single reader"); + }); + + assertEquals(1, queue.size()); + assertTrue(producerPhaser.isTerminated()); + assertEquals(0, batchFailures.get(Entity.TABLE).get()); + } + + @Test + void processBatchQueuesReadResultsAndPromotesFinalBatch() throws Exception { + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + CountDownLatch latch = new CountDownLatch(1); + DefaultRecreateHandler handler = mock(DefaultRecreateHandler.class); + ReindexContext recreateContext = new ReindexContext(); + + recreateContext.add( + Entity.USER, + "user_canonical", + "user_original", + "user_staged", + Set.of("user_existing"), + "user_alias", + List.of("team_alias")); + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.USER, 1))); + setField("taskQueue", queue); + setField("batchSize", new java.util.concurrent.atomic.AtomicReference<>(10)); + setField("config", ReindexingConfiguration.builder().recreateIndex(true).build()); + setField("recreateIndexHandler", handler); + setField("recreateContext", recreateContext); + @SuppressWarnings("unchecked") + Map batchCounters = + (Map) getField("entityBatchCounters"); + @SuppressWarnings("unchecked") + Map batchFailures = + (Map) getField("entityBatchFailures"); + batchCounters.put(Entity.USER, new AtomicInteger(1)); + batchFailures.put(Entity.USER, new AtomicInteger(0)); + + try (MockedConstruction ignored = + mockConstruction( + PaginatedEntitiesSource.class, + (source, context) -> + when(source.readWithCursor(RestUtil.encodeCursor("0"))) + .thenReturn( + (ResultList) new ResultList<>(List.of(mock(EntityInterface.class)))))) { + invokePrivateMethod( + "processBatch", + new Class[] {String.class, int.class, CountDownLatch.class}, + Entity.USER, + 0, + latch); + } + + assertEquals(0, latch.getCount()); + assertEquals(1, queue.size()); + verify(handler).promoteEntityIndex(any(EntityReindexContext.class), eq(true)); + } + + @Test + void handleSinkFailureRoutesProcessAndSinkStagesToRecorder() throws Exception { + IndexingFailureRecorder failureRecorder = mock(IndexingFailureRecorder.class); + setField("failureRecorder", failureRecorder); + + invokePrivateMethod( + "handleSinkFailure", + new Class[] { + String.class, + String.class, + String.class, + String.class, + IndexingFailureRecorder.FailureStage.class + }, + Entity.TABLE, + "1", + "svc.db.table", + "process boom", + IndexingFailureRecorder.FailureStage.PROCESS); + invokePrivateMethod( + "handleSinkFailure", + new Class[] { + String.class, + String.class, + String.class, + String.class, + IndexingFailureRecorder.FailureStage.class + }, + Entity.TABLE, + "2", + "svc.db.table", + "sink boom", + IndexingFailureRecorder.FailureStage.SINK); + + verify(failureRecorder).recordProcessFailure(Entity.TABLE, "1", "svc.db.table", "process boom"); + verify(failureRecorder).recordSinkFailure(Entity.TABLE, "2", "svc.db.table", "sink boom"); + } + + @Test + void isBackpressureActiveTracksQueueFillRatio() throws Exception { + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(10); + ReindexingMetrics metrics = mock(ReindexingMetrics.class); + + for (int i = 0; i < 10; i++) { + queue.add(i); + } + setField("taskQueue", queue); + + try (MockedStatic metricsMock = mockStatic(ReindexingMetrics.class)) { + metricsMock.when(ReindexingMetrics::getInstance).thenReturn(metrics); + + assertTrue((Boolean) invokePrivateMethod("isBackpressureActive", new Class[0])); + verify(metrics).updateQueueFillRatio(100); + + queue.clear(); + assertFalse((Boolean) invokePrivateMethod("isBackpressureActive", new Class[0])); + verify(metrics).updateQueueFillRatio(0); + } + } + + @Test + void calculateNumberOfThreadsHandlesExactRemaindersAndInvalidBatchSize() throws Exception { + assertEquals( + 1, + invokePrivateMethod( + "calculateNumberOfThreads", new Class[] {int.class, int.class}, 10, 0)); + assertEquals( + 2, + invokePrivateMethod( + "calculateNumberOfThreads", new Class[] {int.class, int.class}, 40, 20)); + assertEquals( + 3, + invokePrivateMethod( + "calculateNumberOfThreads", new Class[] {int.class, int.class}, 41, 20)); + } + + @Test + void runConsumerProcessesQueuedWorkUntilPoisonPill() throws Exception { + BulkSink sink = mock(BulkSink.class); + @SuppressWarnings("unchecked") + Source source = mock(Source.class); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + CountDownLatch latch = new CountDownLatch(1); + + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.TABLE, 1))); + setField("config", ReindexingConfiguration.builder().build()); + setField("searchIndexSink", sink); + setField("taskQueue", queue); + when(source.readWithCursor(RestUtil.encodeCursor("0"))) + .thenReturn(new ResultList<>(List.of(mock(EntityInterface.class)))); + + invokePrivateMethod( + "processReadTask", + new Class[] {String.class, Source.class, int.class}, + Entity.TABLE, + source, + 0); + invokePrivateMethod("signalConsumersToStop", new Class[] {int.class}, 1); + invokePrivateMethod("runConsumer", new Class[] {int.class, CountDownLatch.class}, 0, latch); + + verify(sink).write(any(List.class), any(Map.class)); + assertEquals(0, latch.getCount()); + assertEquals(1, executor.getStats().get().getJobStats().getSuccessRecords()); + } + + @Test + void processEntityReindexStopsImmediatelyWhenExecutorIsStopped() throws Exception { + ExecutorService producerExecutor = mock(ExecutorService.class); + ExecutorService jobExecutor = mock(ExecutorService.class); + + setField("producerExecutor", producerExecutor); + setField("jobExecutor", jobExecutor); + ((java.util.concurrent.atomic.AtomicBoolean) getField("stopped")).set(true); + + invokePrivateMethod("processEntityReindex", new Class[] {Set.class}, Set.of(Entity.TABLE)); + + verify(producerExecutor).shutdownNow(); + verify(jobExecutor).shutdownNow(); + ((java.util.concurrent.atomic.AtomicBoolean) getField("stopped")).set(false); + } + + @Test + void cleanupExecutorsShutsDownAllPoolsWhenStillRunning() throws Exception { + ExecutorService consumerExecutor = Executors.newSingleThreadExecutor(); + ExecutorService jobExecutor = Executors.newSingleThreadExecutor(); + ExecutorService producerExecutor = Executors.newSingleThreadExecutor(); + + setField("consumerExecutor", consumerExecutor); + setField("jobExecutor", jobExecutor); + setField("producerExecutor", producerExecutor); + + invokePrivateMethod("cleanupExecutors", new Class[0]); + + assertTrue(consumerExecutor.isShutdown()); + assertTrue(jobExecutor.isShutdown()); + assertTrue(producerExecutor.isShutdown()); + } + + @Test + void removeListenerReturnsExecutorInstance() { + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + + assertSame(executor, executor.addListener(listener).removeListener(listener)); + } + + @Test + void expandEntitiesReturnsIndexedUniverseWhenAllRequested() throws Exception { + when(searchRepository.getEntityIndexMap()) + .thenReturn( + Map.of( + Entity.TABLE, mock(org.openmetadata.search.IndexMapping.class), + Entity.ENTITY_REPORT_DATA, mock(org.openmetadata.search.IndexMapping.class))); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock.when(Entity::getEntityList).thenReturn(Set.of(Entity.TABLE, Entity.USER)); + + @SuppressWarnings("unchecked") + Set expanded = + (Set) + invokePrivateMethod("expandEntities", new Class[] {Set.class}, Set.of("all")); + + assertTrue(expanded.contains(Entity.TABLE)); + assertTrue(expanded.contains(Entity.ENTITY_REPORT_DATA)); + assertFalse(expanded.contains(Entity.USER)); + } + } + + @Test + void calculateThreadConfigurationHonorsConfiguredProducerAndConsumerThreads() throws Exception { + setField( + "config", + ReindexingConfiguration.builder() + .entities(Set.of(Entity.TABLE)) + .producerThreads(6) + .consumerThreads(4) + .build()); + + Object threadConfiguration = + invokePrivateMethod("calculateThreadConfiguration", new Class[] {long.class}, 50_000L); + + assertEquals(6, invokeRecordAccessor(threadConfiguration, "numProducers")); + assertEquals(4, invokeRecordAccessor(threadConfiguration, "numConsumers")); + } + + @Test + void runConsumerContinuesPollingAndExitsWhenInterrupted() throws Exception { + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + CountDownLatch latch = new CountDownLatch(1); + setField("taskQueue", queue); + + Thread consumerThread = + new Thread( + () -> { + try { + invokePrivateMethod( + "runConsumer", new Class[] {int.class, CountDownLatch.class}, 7, latch); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + + consumerThread.start(); + Thread.sleep(250); + consumerThread.interrupt(); + consumerThread.join(2_000); + + assertFalse(consumerThread.isAlive()); + assertEquals(0, latch.getCount()); + } + + @Test + void processTaskRecordsReaderBatchAndHandlesTimeSeriesSinkFailuresWithoutIndexingError() + throws Exception { + BulkSink sink = mock(BulkSink.class); + JobStatsManager statsManager = mock(JobStatsManager.class); + org.openmetadata.service.apps.bundles.searchIndex.stats.EntityStatsTracker tracker = + mock(org.openmetadata.service.apps.bundles.searchIndex.stats.EntityStatsTracker.class); + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + EntityTimeSeriesInterface timeSeriesEntity = mock(EntityTimeSeriesInterface.class); + + executor.addListener(listener); + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.TEST_CASE_RESULT, 1))); + setField("config", ReindexingConfiguration.builder().build()); + setField("searchIndexSink", sink); + setField("statsManager", statsManager); + when(statsManager.getTracker(Entity.TEST_CASE_RESULT)).thenReturn(tracker); + doThrow(new SearchIndexException(new RuntimeException("sink failed"))) + .when(sink) + .write(any(List.class), any(Map.class)); + + invokeProcessTask( + newIndexingTask( + Entity.TEST_CASE_RESULT, + new ResultList<>(List.of(timeSeriesEntity), null, null, 0), + 0)); + + verify(tracker).recordReaderBatch(1, 0, 0); + verify(sink).write(any(List.class), any(Map.class)); + verify(listener) + .onError(eq(Entity.TEST_CASE_RESULT), any(IndexingError.class), any(Stats.class)); + } + + @Test + void processTaskRoutesGenericSinkExceptionsToFailureHandler() throws Exception { + BulkSink sink = mock(BulkSink.class); + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + EntityInterface entity = mock(EntityInterface.class); + + executor.addListener(listener); + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.TABLE, 1))); + setField("config", ReindexingConfiguration.builder().build()); + setField("searchIndexSink", sink); + doThrow(new IllegalStateException("generic sink failure")) + .when(sink) + .write(any(List.class), any(Map.class)); + + invokeProcessTask( + newIndexingTask(Entity.TABLE, new ResultList<>(List.of(entity), null, null, 0), 0)); + + verify(listener).onError(eq(Entity.TABLE), any(IndexingError.class), any(Stats.class)); + } + + @Test + void processEntityTypeUsesTimeSeriesSourcesWithConfiguredWindow() throws Exception { + ExecutorService producerExecutor = mock(ExecutorService.class); + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + Phaser producerPhaser = new Phaser(1); + + doAnswer( + invocation -> { + ((Runnable) invocation.getArgument(0)).run(); + return null; + }) + .when(producerExecutor) + .submit(any(Runnable.class)); + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.TEST_CASE_RESULT, 24))); + setField("producerExecutor", producerExecutor); + setField("taskQueue", queue); + setField( + "config", + ReindexingConfiguration.builder() + .entities(Set.of(Entity.TEST_CASE_RESULT)) + .timeSeriesEntityDays(Map.of(Entity.TEST_CASE_RESULT, 7)) + .build()); + setField("batchSize", new AtomicReference<>(10)); + + try (MockedConstruction ignored = + mockConstruction( + PaginatedEntityTimeSeriesSource.class, + (source, context) -> + when(source.readWithCursor(any())) + .thenReturn( + (ResultList) + new ResultList<>(List.of(mock(EntityTimeSeriesInterface.class)))))) { + invokePrivateMethod( + "processEntityType", + new Class[] {String.class, Phaser.class}, + Entity.TEST_CASE_RESULT, + producerPhaser); + } + + assertFalse(queue.isEmpty()); + assertTrue(producerPhaser.isTerminated()); + } + + @Test + void processEntityTypeDeregistersReaderPartiesWhenSubmissionFails() throws Exception { + ExecutorService producerExecutor = mock(ExecutorService.class); + Phaser producerPhaser = new Phaser(1); + + when(producerExecutor.submit(any(Runnable.class))) + .thenThrow(new IllegalStateException("submit failed")); + executor.getStats().set(statsWithEntityTotals(Map.of(Entity.USER, 40))); + setField("producerExecutor", producerExecutor); + setField("taskQueue", new LinkedBlockingQueue<>()); + setField("config", ReindexingConfiguration.builder().entities(Set.of(Entity.USER)).build()); + setField("batchSize", new AtomicReference<>(10)); + + invokePrivateMethod( + "processEntityType", + new Class[] {String.class, Phaser.class}, + Entity.USER, + producerPhaser); + + assertTrue(producerPhaser.isTerminated()); + } + + @Test + void processKeysetBatchesStopsWhenReaderReachesEndCursorBoundary() throws Exception { + LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + Phaser producerPhaser = new Phaser(1); + String boundaryCursor = "{\"name\":\"orders\",\"id\":\"2\"}"; + String endCursor = RestUtil.encodeCursor(boundaryCursor); + ResultList page = new ResultList<>(List.of("entity"), null, null, boundaryCursor, 1); + + setField("taskQueue", queue); + @SuppressWarnings("unchecked") + Map batchCounters = + (Map) getField("entityBatchCounters"); + @SuppressWarnings("unchecked") + Map batchFailures = + (Map) getField("entityBatchFailures"); + batchCounters.put(Entity.TABLE, new AtomicInteger(1)); + batchFailures.put(Entity.TABLE, new AtomicInteger(0)); + + invokePrivateMethod( + "processKeysetBatches", + new Class[] { + String.class, + int.class, + int.class, + String.class, + SearchIndexExecutor.KeysetBatchReader.class, + Phaser.class, + String.class + }, + Entity.TABLE, + 10, + 5, + null, + (SearchIndexExecutor.KeysetBatchReader) cursor -> page, + producerPhaser, + endCursor); + + assertEquals(1, queue.size()); + assertEquals(0, batchFailures.get(Entity.TABLE).get()); + assertTrue(producerPhaser.isTerminated()); + } + + @Test + void processKeysetBatchesMarksFailuresForUnexpectedExceptions() throws Exception { + Phaser producerPhaser = new Phaser(1); + @SuppressWarnings("unchecked") + Map batchCounters = + (Map) getField("entityBatchCounters"); + @SuppressWarnings("unchecked") + Map batchFailures = + (Map) getField("entityBatchFailures"); + batchCounters.put(Entity.TABLE, new AtomicInteger(1)); + batchFailures.put(Entity.TABLE, new AtomicInteger(0)); + setField("taskQueue", new LinkedBlockingQueue<>()); + + invokePrivateMethod( + "processKeysetBatches", + new Class[] { + String.class, + int.class, + int.class, + String.class, + SearchIndexExecutor.KeysetBatchReader.class, + Phaser.class + }, + Entity.TABLE, + 5, + 5, + null, + (SearchIndexExecutor.KeysetBatchReader) + cursor -> { + throw new IllegalStateException("unexpected"); + }, + producerPhaser); + + assertEquals(1, batchFailures.get(Entity.TABLE).get()); + assertTrue(producerPhaser.isTerminated()); + } + + /** + * Validates the full cursor decode → BoundedListFilter flow: an encoded boundary cursor + * is decoded and used to construct a filter with the correct SQL boundary condition. + * This is the core mechanism that replaces the broken Java-side hasReachedEndCursor comparison. + */ + @Test + @SuppressWarnings("unchecked") + void encodedBoundaryCursorProducesCorrectBoundedFilter() { + // The exact cursor that would be produced by getCursorAtOffset for entity "Foxtrot" + String boundaryCursorJson = + "{\"name\":\"Foxtrot\",\"id\":\"00000000-0000-0000-0000-000000000006\"}"; + String encodedBoundary = RestUtil.encodeCursor(boundaryCursorJson); + + // Decode — same logic as submitEntityReaders + String decoded = RestUtil.decodeCursor(encodedBoundary); + Map cursorMap = + org.openmetadata.schema.utils.JsonUtils.readValue(decoded, Map.class); + + assertEquals("Foxtrot", cursorMap.get("name")); + assertEquals("00000000-0000-0000-0000-000000000006", cursorMap.get("id")); + + // Construct BoundedListFilter with decoded values + org.openmetadata.service.jdbi3.BoundedListFilter filter = + new org.openmetadata.service.jdbi3.BoundedListFilter( + org.openmetadata.schema.type.Include.ALL, cursorMap.get("name"), cursorMap.get("id")); + + String condition = filter.getCondition(null); + assertTrue(condition.contains("name < :reindexEndName")); + assertTrue(condition.contains("name = :reindexEndName AND id <= :reindexEndId")); + assertEquals("Foxtrot", filter.getQueryParams().get("reindexEndName")); + assertEquals( + "00000000-0000-0000-0000-000000000006", filter.getQueryParams().get("reindexEndId")); + } + + /** + * Verifies that a BoundedListFilter and a plain ListFilter produce different conditions, + * confirming the non-last reader gets a bounded query while the last reader does not. + */ + @Test + void boundedVsUnboundedFilterProduceDifferentConditions() { + ListFilter unbounded = new ListFilter(org.openmetadata.schema.type.Include.ALL); + org.openmetadata.service.jdbi3.BoundedListFilter bounded = + new org.openmetadata.service.jdbi3.BoundedListFilter( + org.openmetadata.schema.type.Include.ALL, + "Foxtrot", + "00000000-0000-0000-0000-000000000006"); + + String unboundedCond = unbounded.getCondition(null); + String boundedCond = bounded.getCondition(null); + + assertFalse(unboundedCond.contains("reindexEndName")); + assertTrue(boundedCond.contains("reindexEndName")); + assertTrue(boundedCond.startsWith(unboundedCond)); + } + + /** + * Validates that the old Java-side cursor comparison no longer applies to entity cursors. + * This is the exact scenario that caused the bug: "echo".compareTo("Foxtrot") > 0 in Java + * but "echo" < "Foxtrot" in MySQL case-insensitive collation. + */ + @Test + void hasReachedEndCursorNoLongerComparesEntityCursors() throws Exception { + // This is the exact pair that triggered the bug: + // Java: "echo" > "Foxtrot" (e=101 > F=70) → old code returned TRUE (stop reader) + // MySQL: "echo" < "Foxtrot" (case-insensitive: e < f) → reader should continue + String echoCursor = + RestUtil.encodeCursor( + "{\"name\":\"echo\",\"id\":\"00000000-0000-0000-0000-000000000005\"}"); + String foxtrotCursor = + RestUtil.encodeCursor( + "{\"name\":\"Foxtrot\",\"id\":\"00000000-0000-0000-0000-000000000006\"}"); + + // After fix: hasReachedEndCursor returns FALSE for entity cursors (boundary is in SQL now) + assertFalse( + (Boolean) + invokePrivateMethod( + "hasReachedEndCursor", + new Class[] {String.class, String.class}, + echoCursor, + foxtrotCursor), + "Entity cursor comparison must not happen in Java — SQL boundary handles it"); + } + + /** + * Verifies that the old bug scenario is now impossible: mixed-case names at boundaries + * cannot cause missing entities because the boundary is enforced in SQL, not Java. + */ + @Test + void mixedCaseEntityNamesAtBoundaryProduceBoundedSqlCondition() { + // Simulate the exact scenario: boundary entity is "Foxtrot" + org.openmetadata.service.jdbi3.BoundedListFilter filter = + new org.openmetadata.service.jdbi3.BoundedListFilter( + org.openmetadata.schema.type.Include.ALL, + "Foxtrot", + "00000000-0000-0000-0000-000000000006"); + + String condition = filter.getCondition(null); + + // The SQL condition ensures the DB collation handles the comparison. + // On MySQL: WHERE ... AND (name < 'Foxtrot' OR (name = 'Foxtrot' AND id <= 'uuid')) + // The DB evaluates "echo" < "Foxtrot" as TRUE (case-insensitive), so "echo" IS included. + // "Foxtrot" itself is included (id <= boundary id). + // "golf" is excluded (name > "Foxtrot" case-insensitively). + assertTrue(condition.contains("name < :reindexEndName")); + assertTrue(condition.contains("name = :reindexEndName AND id <= :reindexEndId")); + assertEquals("Foxtrot", filter.getQueryParams().get("reindexEndName")); + } + + private Stats initializeStats(Set entities) { + Stats stats = executor.initializeTotalRecords(entities); + if (stats.getEntityStats() == null) { + stats.setEntityStats(new EntityStats()); + } + return stats; + } + + private Stats statsWithEntityTotals(Map entityTotals) { + Stats stats = new Stats(); + EntityStats entityStats = new EntityStats(); + int totalRecords = 0; + + for (Map.Entry entry : entityTotals.entrySet()) { + totalRecords += entry.getValue(); + entityStats + .getAdditionalProperties() + .put( + entry.getKey(), + new StepStats() + .withTotalRecords(entry.getValue()) + .withSuccessRecords(0) + .withFailedRecords(0)); + } + + stats.setEntityStats(entityStats); + stats.setJobStats( + new StepStats().withTotalRecords(totalRecords).withSuccessRecords(0).withFailedRecords(0)); + stats.setReaderStats( + new StepStats() + .withTotalRecords(totalRecords) + .withSuccessRecords(0) + .withFailedRecords(0) + .withWarningRecords(0)); + stats.setSinkStats( + new StepStats().withTotalRecords(0).withSuccessRecords(0).withFailedRecords(0)); + stats.setProcessStats( + new StepStats().withTotalRecords(0).withSuccessRecords(0).withFailedRecords(0)); + return stats; + } + + private Object invokePrivateMethod(String methodName, Class[] parameterTypes, Object... args) + throws Exception { + Method method = SearchIndexExecutor.class.getDeclaredMethod(methodName, parameterTypes); + method.setAccessible(true); + return method.invoke(executor, args); + } + + private void setField(String fieldName, Object value) throws Exception { + Field field = SearchIndexExecutor.class.getDeclaredField(fieldName); + field.setAccessible(true); + field.set(executor, value); + } + + private Object getField(String fieldName) throws Exception { + Field field = SearchIndexExecutor.class.getDeclaredField(fieldName); + field.setAccessible(true); + return field.get(executor); + } + + private Object newIndexingTask(String entityType, ResultList entities, int offset) + throws Exception { + Class taskClass = + Class.forName( + "org.openmetadata.service.apps.bundles.searchIndex.SearchIndexExecutor$IndexingTask"); + var constructor = taskClass.getDeclaredConstructor(String.class, ResultList.class, int.class); + constructor.setAccessible(true); + return constructor.newInstance(entityType, entities, offset); + } + + private void invokeProcessTask(Object task) throws Exception { + Method method = SearchIndexExecutor.class.getDeclaredMethod("processTask", task.getClass()); + method.setAccessible(true); + method.invoke(executor, task); + } + + private Object invokeRecordAccessor(Object record, String accessor) throws Exception { + Method method = record.getClass().getDeclaredMethod(accessor); + method.setAccessible(true); + return method.invoke(record); + } + + private Object invokeTaskAccessor(Object task, String accessor) throws Exception { + Method method = task.getClass().getDeclaredMethod(accessor); + method.setAccessible(true); + return method.invoke(task); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexFailureScenarioTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexFailureScenarioTest.java new file mode 100644 index 000000000000..950fcde70c9c --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexFailureScenarioTest.java @@ -0,0 +1,522 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; + +import es.co.elastic.clients.elasticsearch.ElasticsearchClient; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicLong; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.search.IndexMapping; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.search.elasticsearch.ElasticSearchClient; + +/** + * Comprehensive tests for SearchIndex stats accuracy across all failure scenarios: + * 1. Request entity too large (413) from ES/OS + * 2. Entity read failures + * 3. Entity build failures + * 4. Partial bulk failures + * 5. Complete bulk request failures + * 6. Reader exceptions + * 7. Sink exceptions + */ +@ExtendWith(MockitoExtension.class) +class SearchIndexFailureScenarioTest { + + @Mock private SearchRepository searchRepository; + @Mock private ElasticSearchClient searchClient; + @Mock private ElasticsearchClient restHighLevelClient; + @Mock private IndexMapping indexMapping; + @Mock private CollectionDAO collectionDAO; + + @BeforeEach + void setUp() { + lenient().when(searchRepository.getSearchClient()).thenReturn(searchClient); + lenient().when(searchClient.getNewClient()).thenReturn(restHighLevelClient); + lenient().when(searchRepository.getClusterAlias()).thenReturn("default"); + lenient().when(indexMapping.getIndexName("default")).thenReturn("test_index"); + lenient().when(searchRepository.getIndexMapping(anyString())).thenReturn(indexMapping); + } + + @Nested + @DisplayName("Scenario 1: Request Entity Too Large (413)") + class RequestEntityTooLargeTests { + + @Test + @DisplayName("Should detect 413 error as payload too large") + void testDetect413Error() throws Exception { + ElasticSearchBulkSink.CustomBulkProcessor processor = + getCustomBulkProcessor(new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L)); + + assertTrue(invokeIsPayloadTooLargeError(processor, "Request entity too large")); + assertTrue(invokeIsPayloadTooLargeError(processor, "HTTP/1.1 413 Payload Too Large")); + assertTrue(invokeIsPayloadTooLargeError(processor, "content too long")); + assertTrue(invokeIsPayloadTooLargeError(processor, "Error code: 413")); + } + + @Test + @DisplayName("Should detect 413 error as backpressure trigger") + void test413TriggersBackpressure() throws Exception { + ElasticSearchBulkSink.CustomBulkProcessor processor = + getCustomBulkProcessor(new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L)); + + assertTrue(invokeShouldRetry(processor, 0, "Request entity too large")); + assertTrue(invokeShouldRetry(processor, 0, "413")); + } + + @Test + @DisplayName("BulkSink should identify 413 as retryable error") + void testBulkSinkRetries413() throws Exception { + ElasticSearchBulkSink sink = new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L); + ElasticSearchBulkSink.CustomBulkProcessor processor = getCustomBulkProcessor(sink); + + assertTrue(invokeShouldRetry(processor, 0, "Request entity too large")); + assertTrue(invokeShouldRetry(processor, 0, "Content too long")); + assertTrue(invokeShouldRetry(processor, 0, "413")); + assertFalse(invokeShouldRetry(processor, 5, "Request entity too large")); + } + } + + @Nested + @DisplayName("Scenario 2: Entity Read Failures") + class EntityReadFailureTests { + + @Test + @DisplayName("Reader failures should update reader stats") + void testReaderFailuresUpdateStats() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateReaderStats(0, 10, 0); + + Stats updatedStats = executor.getStats().get(); + assertNotNull(updatedStats); + assertEquals(0, updatedStats.getReaderStats().getSuccessRecords()); + assertEquals(10, updatedStats.getReaderStats().getFailedRecords()); + } + + @Test + @DisplayName("Partial read failures should be tracked correctly") + void testPartialReadFailures() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateReaderStats(90, 10, 0); + executor.updateReaderStats(85, 15, 0); + + Stats updatedStats = executor.getStats().get(); + assertEquals(175, updatedStats.getReaderStats().getSuccessRecords()); + assertEquals(25, updatedStats.getReaderStats().getFailedRecords()); + } + } + + @Nested + @DisplayName("Scenario 3: Entity Build Failures") + class EntityBuildFailureTests { + + @Test + @DisplayName("Process failures should be tracked in totalFailed") + void testProcessFailuresTracked() throws Exception { + ElasticSearchBulkSink sink = new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L); + + // Failures during entity processing (building search docs) are tracked in totalFailed + Field totalFailedField = ElasticSearchBulkSink.class.getDeclaredField("totalFailed"); + totalFailedField.setAccessible(true); + AtomicLong totalFailed = (AtomicLong) totalFailedField.get(sink); + totalFailed.set(5); + + Method updateStatsMethod = ElasticSearchBulkSink.class.getDeclaredMethod("updateStats"); + updateStatsMethod.setAccessible(true); + updateStatsMethod.invoke(sink); + + StepStats stats = sink.getStats(); + assertEquals(5, stats.getFailedRecords()); + } + } + + @Nested + @DisplayName("Scenario 4: Partial Bulk Failures") + class PartialBulkFailureTests { + + @Test + @DisplayName("Partial bulk failures should correctly split success and failure counts") + void testPartialBulkFailureStats() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + StepStats batchStats = new StepStats().withSuccessRecords(8).withFailedRecords(2); + executor.updateStats("table", batchStats); + + Stats finalStats = executor.getStats().get(); + StepStats entityStats = finalStats.getEntityStats().getAdditionalProperties().get("table"); + + assertEquals(8, entityStats.getSuccessRecords()); + assertEquals(2, entityStats.getFailedRecords()); + assertEquals(8, finalStats.getJobStats().getSuccessRecords()); + assertEquals(2, finalStats.getJobStats().getFailedRecords()); + } + } + + @Nested + @DisplayName("Scenario 5: Complete Bulk Request Failures") + class CompleteBulkFailureTests { + + @Test + @DisplayName("Complete bulk failure should mark all records as failed") + void testCompleteBulkFailure() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + StepStats batchStats = new StepStats().withSuccessRecords(0).withFailedRecords(100); + executor.updateStats("table", batchStats); + + Stats finalStats = executor.getStats().get(); + assertEquals(0, finalStats.getJobStats().getSuccessRecords()); + assertEquals(100, finalStats.getJobStats().getFailedRecords()); + } + } + + @Nested + @DisplayName("Scenario 6: Stats Consistency") + class StatsConsistencyTests { + + @Test + @DisplayName("Total should equal success + failed after all operations") + void testTotalEqualsSuccessPlusFailed() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table", "dashboard"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn( + Map.of("table", mock(IndexMapping.class), "dashboard", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + stats.getEntityStats().getAdditionalProperties().get("table").setTotalRecords(100); + stats.getEntityStats().getAdditionalProperties().get("dashboard").setTotalRecords(50); + stats.getJobStats().setTotalRecords(150); + stats.getReaderStats().setTotalRecords(150); + executor.getStats().set(stats); + + executor.updateStats("table", new StepStats().withSuccessRecords(90).withFailedRecords(10)); + executor.updateStats( + "dashboard", new StepStats().withSuccessRecords(45).withFailedRecords(5)); + executor.updateReaderStats(135, 15, 0); + + Stats finalStats = executor.getStats().get(); + + int jobSuccess = finalStats.getJobStats().getSuccessRecords(); + int jobFailed = finalStats.getJobStats().getFailedRecords(); + int jobTotal = finalStats.getJobStats().getTotalRecords(); + + assertEquals(135, jobSuccess); + assertEquals(15, jobFailed); + assertEquals(jobSuccess + jobFailed, jobTotal); + } + + @Test + @DisplayName("Entity stats sum should equal job stats") + void testEntityStatsSumEqualsJobStats() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table", "dashboard", "pipeline"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn( + Map.of( + "table", mock(IndexMapping.class), + "dashboard", mock(IndexMapping.class), + "pipeline", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateStats("table", new StepStats().withSuccessRecords(50).withFailedRecords(5)); + executor.updateStats( + "dashboard", new StepStats().withSuccessRecords(30).withFailedRecords(3)); + executor.updateStats("pipeline", new StepStats().withSuccessRecords(20).withFailedRecords(2)); + + Stats finalStats = executor.getStats().get(); + + int entitySuccessSum = 0; + int entityFailedSum = 0; + for (StepStats entityStats : finalStats.getEntityStats().getAdditionalProperties().values()) { + entitySuccessSum += entityStats.getSuccessRecords(); + entityFailedSum += entityStats.getFailedRecords(); + } + + assertEquals(entitySuccessSum, finalStats.getJobStats().getSuccessRecords()); + assertEquals(entityFailedSum, finalStats.getJobStats().getFailedRecords()); + } + } + + @Nested + @DisplayName("Scenario 7: Error Type Detection") + class ErrorTypeDetectionTests { + + @Test + @DisplayName("Should correctly identify all retryable error types") + void testAllRetryableErrorTypes() throws Exception { + ElasticSearchBulkSink sink = new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L); + + Field field = ElasticSearchBulkSink.class.getDeclaredField("bulkProcessor"); + field.setAccessible(true); + ElasticSearchBulkSink.CustomBulkProcessor processor = + (ElasticSearchBulkSink.CustomBulkProcessor) field.get(sink); + + Method method = + ElasticSearchBulkSink.CustomBulkProcessor.class.getDeclaredMethod( + "shouldRetry", int.class, Throwable.class); + method.setAccessible(true); + + String[] retryableErrors = { + "rejected_execution_exception", + "EsRejectedExecutionException", + "RemoteTransportException", + "ConnectException", + "timeout", + "Request entity too large", + "Content too long", + "413", + "circuit_breaking_exception", + "too_many_requests" + }; + + for (String errorMessage : retryableErrors) { + assertTrue( + (boolean) method.invoke(processor, 0, new RuntimeException(errorMessage)), + "Should retry for: " + errorMessage); + } + } + + @Test + @DisplayName("Should NOT retry non-retryable errors") + void testNonRetryableErrors() throws Exception { + ElasticSearchBulkSink sink = new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L); + + Field field = ElasticSearchBulkSink.class.getDeclaredField("bulkProcessor"); + field.setAccessible(true); + ElasticSearchBulkSink.CustomBulkProcessor processor = + (ElasticSearchBulkSink.CustomBulkProcessor) field.get(sink); + + Method method = + ElasticSearchBulkSink.CustomBulkProcessor.class.getDeclaredMethod( + "shouldRetry", int.class, Throwable.class); + method.setAccessible(true); + + String[] nonRetryableErrors = { + "index_not_found_exception", + "mapper_parsing_exception", + "document_parsing_exception", + "invalid_type_name_exception" + }; + + for (String errorMessage : nonRetryableErrors) { + assertFalse( + (boolean) method.invoke(processor, 0, new RuntimeException(errorMessage)), + "Should NOT retry for: " + errorMessage); + } + } + + @Test + @DisplayName("Should correctly identify backpressure errors") + void testBackpressureErrorDetection() throws Exception { + ElasticSearchBulkSink.CustomBulkProcessor processor = + getCustomBulkProcessor(new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L)); + + String[] backpressureErrors = { + "rejected_execution_exception", + "circuit_breaking_exception", + "too_many_requests", + "Request entity too large", + "Content too long", + "413" + }; + + for (String errorMessage : backpressureErrors) { + assertTrue( + invokeShouldRetry(processor, 0, errorMessage), + "Should be backpressure for: " + errorMessage); + } + } + } + + @Nested + @DisplayName("Scenario 8: Multi-Batch Stats Accumulation") + class MultiBatchStatsAccumulationTests { + + @Test + @DisplayName("Stats should accumulate correctly across multiple batches") + void testMultiBatchAccumulation() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + for (int i = 0; i < 10; i++) { + executor.updateStats("table", new StepStats().withSuccessRecords(9).withFailedRecords(1)); + executor.updateReaderStats(10, 0, 0); + executor.updateSinkTotalSubmitted(10); + } + + Stats finalStats = executor.getStats().get(); + + assertEquals(90, finalStats.getJobStats().getSuccessRecords()); + assertEquals(10, finalStats.getJobStats().getFailedRecords()); + assertEquals(100, finalStats.getReaderStats().getSuccessRecords()); + assertEquals(0, finalStats.getReaderStats().getFailedRecords()); + assertEquals(100, finalStats.getSinkStats().getTotalRecords()); + } + + @Test + @DisplayName("Interleaved success and failure batches should accumulate correctly") + void testInterleavedSuccessAndFailure() { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateStats("table", new StepStats().withSuccessRecords(100).withFailedRecords(0)); + executor.updateStats("table", new StepStats().withSuccessRecords(0).withFailedRecords(50)); + executor.updateStats("table", new StepStats().withSuccessRecords(75).withFailedRecords(25)); + + Stats finalStats = executor.getStats().get(); + + assertEquals(175, finalStats.getJobStats().getSuccessRecords()); + assertEquals(75, finalStats.getJobStats().getFailedRecords()); + } + } + + @Nested + @DisplayName("Scenario 9: Concurrent Stats Updates") + class ConcurrentStatsUpdateTests { + + @Test + @DisplayName("Concurrent updates should not lose data") + void testConcurrentUpdates() throws Exception { + SearchIndexExecutor executor = new SearchIndexExecutor(collectionDAO, searchRepository); + + Set entities = Set.of("table"); + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + int threadCount = 10; + int updatesPerThread = 100; + Thread[] threads = new Thread[threadCount]; + + for (int i = 0; i < threadCount; i++) { + threads[i] = + new Thread( + () -> { + for (int j = 0; j < updatesPerThread; j++) { + executor.updateStats( + "table", new StepStats().withSuccessRecords(1).withFailedRecords(0)); + } + }); + } + + for (Thread thread : threads) { + thread.start(); + } + + for (Thread thread : threads) { + thread.join(); + } + + Stats finalStats = executor.getStats().get(); + int expectedTotal = threadCount * updatesPerThread; + + assertEquals(expectedTotal, finalStats.getJobStats().getSuccessRecords()); + } + } + + private ElasticSearchBulkSink.CustomBulkProcessor getCustomBulkProcessor( + ElasticSearchBulkSink sink) throws Exception { + Field field = ElasticSearchBulkSink.class.getDeclaredField("bulkProcessor"); + field.setAccessible(true); + return (ElasticSearchBulkSink.CustomBulkProcessor) field.get(sink); + } + + private boolean invokeShouldRetry( + ElasticSearchBulkSink.CustomBulkProcessor processor, int attemptNumber, String errorMessage) + throws Exception { + Method method = + ElasticSearchBulkSink.CustomBulkProcessor.class.getDeclaredMethod( + "shouldRetry", int.class, Throwable.class); + method.setAccessible(true); + Throwable error = + errorMessage == null ? new RuntimeException() : new RuntimeException(errorMessage); + return (boolean) method.invoke(processor, attemptNumber, error); + } + + private boolean invokeIsPayloadTooLargeError( + ElasticSearchBulkSink.CustomBulkProcessor processor, String errorMessage) throws Exception { + Method method = + ElasticSearchBulkSink.CustomBulkProcessor.class.getDeclaredMethod( + "isPayloadTooLargeError", Throwable.class); + method.setAccessible(true); + Throwable error = + errorMessage == null ? new RuntimeException() : new RuntimeException(errorMessage); + return (boolean) method.invoke(processor, error); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexStatsTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexStatsTest.java new file mode 100644 index 000000000000..6331326f6ed3 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SearchIndexStatsTest.java @@ -0,0 +1,444 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; + +import es.co.elastic.clients.elasticsearch.ElasticsearchClient; +import java.lang.reflect.Method; +import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.schema.system.StepStats; +import org.openmetadata.search.IndexMapping; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.search.SearchRepository; +import org.openmetadata.service.search.elasticsearch.ElasticSearchClient; + +@ExtendWith(MockitoExtension.class) +class SearchIndexStatsTest { + + @Mock private SearchRepository searchRepository; + @Mock private ElasticSearchClient searchClient; + @Mock private ElasticsearchClient restHighLevelClient; + @Mock private IndexMapping indexMapping; + @Mock private CollectionDAO collectionDAO; + + @BeforeEach + void setUp() { + lenient().when(searchRepository.getSearchClient()).thenReturn(searchClient); + lenient().when(searchClient.getNewClient()).thenReturn(restHighLevelClient); + lenient().when(searchRepository.getClusterAlias()).thenReturn("default"); + lenient().when(indexMapping.getIndexName("default")).thenReturn("test_index"); + lenient().when(searchRepository.getIndexMapping(anyString())).thenReturn(indexMapping); + } + + @Nested + @DisplayName("BulkSink Stats Tests") + class BulkSinkStatsTests { + + private ElasticSearchBulkSink elasticSearchBulkSink; + + @BeforeEach + void setUp() { + elasticSearchBulkSink = new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L); + } + + @Test + @DisplayName("Initial stats should be zero") + void testInitialStatsAreZero() { + StepStats stats = elasticSearchBulkSink.getStats(); + assertNotNull(stats); + assertEquals(0, stats.getTotalRecords()); + assertEquals(0, stats.getSuccessRecords()); + assertEquals(0, stats.getFailedRecords()); + } + } + + @Nested + @DisplayName("Retry Logic Tests") + class RetryLogicTests { + + @Test + @DisplayName("Should identify 'Request entity too large' as retryable error") + void testRequestEntityTooLargeIsRetryable() throws Exception { + ElasticSearchBulkSink sink = new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L); + + ElasticSearchBulkSink.CustomBulkProcessor processor = getCustomBulkProcessor(sink); + + assertTrue(invokeIsPayloadTooLargeError(processor, "Request entity too large")); + assertTrue(invokeIsPayloadTooLargeError(processor, "Content too long")); + assertTrue(invokeIsPayloadTooLargeError(processor, "HTTP 413 error")); + } + } + + @Nested + @DisplayName("SearchIndexExecutor Stats Tests") + class ExecutorStatsTests { + + private SearchIndexExecutor executor; + + @BeforeEach + void setUp() { + executor = new SearchIndexExecutor(collectionDAO, searchRepository); + } + + @Test + @DisplayName("Stats initialization should set all values correctly") + void testStatsInitialization() { + Set entities = Set.of("table", "dashboard"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn( + Map.of("table", mock(IndexMapping.class), "dashboard", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + + assertNotNull(stats); + assertNotNull(stats.getJobStats()); + assertNotNull(stats.getReaderStats()); + assertNotNull(stats.getSinkStats()); + assertNotNull(stats.getEntityStats()); + + assertEquals(0, stats.getJobStats().getSuccessRecords()); + assertEquals(0, stats.getJobStats().getFailedRecords()); + assertEquals(0, stats.getReaderStats().getSuccessRecords()); + assertEquals(0, stats.getReaderStats().getFailedRecords()); + assertEquals(0, stats.getSinkStats().getSuccessRecords()); + assertEquals(0, stats.getSinkStats().getFailedRecords()); + } + + @Test + @DisplayName("updateStats should correctly accumulate values") + void testUpdateStatsAccumulation() { + Set entities = Set.of("table"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + StepStats batchStats = new StepStats().withSuccessRecords(5).withFailedRecords(2); + executor.updateStats("table", batchStats); + + Stats updatedStats = executor.getStats().get(); + assertNotNull(updatedStats); + + StepStats entityStats = updatedStats.getEntityStats().getAdditionalProperties().get("table"); + assertNotNull(entityStats); + assertEquals(5, entityStats.getSuccessRecords()); + assertEquals(2, entityStats.getFailedRecords()); + + assertEquals(5, updatedStats.getJobStats().getSuccessRecords()); + assertEquals(2, updatedStats.getJobStats().getFailedRecords()); + } + + @Test + @DisplayName("updateReaderStats should correctly track reader operations") + void testUpdateReaderStats() { + Set entities = Set.of("table"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateReaderStats(10, 2, 0); + + Stats updatedStats = executor.getStats().get(); + assertNotNull(updatedStats); + assertEquals(10, updatedStats.getReaderStats().getSuccessRecords()); + assertEquals(2, updatedStats.getReaderStats().getFailedRecords()); + + executor.updateReaderStats(5, 1, 0); + + updatedStats = executor.getStats().get(); + assertEquals(15, updatedStats.getReaderStats().getSuccessRecords()); + assertEquals(3, updatedStats.getReaderStats().getFailedRecords()); + } + + @Test + @DisplayName("updateSinkTotalSubmitted should correctly track submitted records") + void testUpdateSinkTotalSubmitted() { + Set entities = Set.of("table"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateSinkTotalSubmitted(10); + + Stats updatedStats = executor.getStats().get(); + assertNotNull(updatedStats); + assertEquals(10, updatedStats.getSinkStats().getTotalRecords()); + + executor.updateSinkTotalSubmitted(5); + + updatedStats = executor.getStats().get(); + assertEquals(15, updatedStats.getSinkStats().getTotalRecords()); + } + } + + @Nested + @DisplayName("Backpressure Detection Tests") + class BackpressureDetectionTests { + + @Test + @DisplayName("Should detect payload-too-large errors as retryable backpressure") + void testPayloadTooLargeDetectedAsBackpressure() throws Exception { + ElasticSearchBulkSink.CustomBulkProcessor processor = + getCustomBulkProcessor(new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L)); + + assertTrue(invokeShouldRetry(processor, 0, "Request entity too large")); + assertTrue(invokeShouldRetry(processor, 0, "Content too long for bulk request")); + assertTrue(invokeShouldRetry(processor, 0, "HTTP 413: Payload too large")); + } + + @Test + @DisplayName("Should detect rejected_execution_exception as backpressure error") + void testRejectedExecutionDetectedAsBackpressure() throws Exception { + ElasticSearchBulkSink.CustomBulkProcessor processor = + getCustomBulkProcessor(new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L)); + + assertTrue(invokeShouldRetry(processor, 0, "rejected_execution_exception")); + assertTrue(invokeShouldRetry(processor, 0, "circuit_breaking_exception")); + assertTrue(invokeShouldRetry(processor, 0, "too_many_requests")); + } + + @Test + @DisplayName( + "Should detect only known backpressure errors while treating null messages as retryable") + void testNormalErrorsNotBackpressure() throws Exception { + ElasticSearchBulkSink.CustomBulkProcessor processor = + getCustomBulkProcessor(new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L)); + + assertFalse(invokeShouldRetry(processor, 0, "Index not found")); + assertFalse(invokeShouldRetry(processor, 0, "Document parsing exception")); + assertFalse(invokeShouldRetry(processor, 0, "Mapping error")); + assertTrue(invokeShouldRetry(processor, 0, null)); + } + + @Test + @DisplayName("Should identify payload too large error correctly") + void testIsPayloadTooLargeError() throws Exception { + ElasticSearchBulkSink.CustomBulkProcessor processor = + getCustomBulkProcessor(new ElasticSearchBulkSink(searchRepository, 10, 2, 1000000L)); + + assertTrue(invokeIsPayloadTooLargeError(processor, "Request entity too large")); + assertTrue(invokeIsPayloadTooLargeError(processor, "Content too long")); + assertTrue(invokeIsPayloadTooLargeError(processor, "error code: 413")); + + assertFalse(invokeIsPayloadTooLargeError(processor, "rejected_execution_exception")); + assertFalse(invokeIsPayloadTooLargeError(processor, "timeout")); + assertFalse(invokeIsPayloadTooLargeError(processor, null)); + } + } + + @Nested + @DisplayName("Stats Consistency Tests") + class StatsConsistencyTests { + + private SearchIndexExecutor executor; + + @BeforeEach + void setUp() { + executor = new SearchIndexExecutor(collectionDAO, searchRepository); + } + + @Test + @DisplayName("Job stats should match sum of entity stats") + void testJobStatsMatchEntityStats() { + Set entities = Set.of("table", "dashboard", "pipeline"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn( + Map.of( + "table", mock(IndexMapping.class), + "dashboard", mock(IndexMapping.class), + "pipeline", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateStats("table", new StepStats().withSuccessRecords(10).withFailedRecords(2)); + executor.updateStats("dashboard", new StepStats().withSuccessRecords(5).withFailedRecords(1)); + executor.updateStats("pipeline", new StepStats().withSuccessRecords(8).withFailedRecords(3)); + + Stats finalStats = executor.getStats().get(); + + int expectedSuccess = 10 + 5 + 8; + int expectedFailed = 2 + 1 + 3; + + assertEquals(expectedSuccess, finalStats.getJobStats().getSuccessRecords()); + assertEquals(expectedFailed, finalStats.getJobStats().getFailedRecords()); + } + + @Test + @DisplayName("Multiple updates to same entity should accumulate correctly") + void testMultipleUpdatesToSameEntity() { + Set entities = Set.of("table"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + executor.updateStats("table", new StepStats().withSuccessRecords(10).withFailedRecords(2)); + executor.updateStats("table", new StepStats().withSuccessRecords(5).withFailedRecords(1)); + executor.updateStats("table", new StepStats().withSuccessRecords(3).withFailedRecords(0)); + + Stats finalStats = executor.getStats().get(); + StepStats tableStats = finalStats.getEntityStats().getAdditionalProperties().get("table"); + + assertEquals(18, tableStats.getSuccessRecords()); + assertEquals(3, tableStats.getFailedRecords()); + + assertEquals(18, finalStats.getJobStats().getSuccessRecords()); + assertEquals(3, finalStats.getJobStats().getFailedRecords()); + } + + @Test + @DisplayName("Stats should handle null stats object gracefully") + void testNullStatsHandling() { + executor.updateStats("table", new StepStats().withSuccessRecords(10).withFailedRecords(2)); + executor.updateReaderStats(5, 1, 0); + executor.updateSinkTotalSubmitted(10); + } + + @Test + @DisplayName("Entity total should be adjusted when success + failed exceeds initial total") + void testEntityTotalAdjustedWhenExceeded() { + Set entities = Set.of("table"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + // Initial total is 0 (mocked). Simulate batches that exceed it. + executor.updateStats("table", new StepStats().withSuccessRecords(50).withFailedRecords(2)); + executor.updateStats("table", new StepStats().withSuccessRecords(55).withFailedRecords(1)); + + Stats finalStats = executor.getStats().get(); + StepStats tableStats = finalStats.getEntityStats().getAdditionalProperties().get("table"); + + assertEquals(105, tableStats.getSuccessRecords()); + assertEquals(3, tableStats.getFailedRecords()); + // Total should have been bumped to success + failed + assertEquals(108, tableStats.getTotalRecords()); + + // Job total should also reflect the adjusted entity total + assertEquals(108, finalStats.getJobStats().getTotalRecords()); + assertEquals(105, finalStats.getJobStats().getSuccessRecords()); + assertEquals(3, finalStats.getJobStats().getFailedRecords()); + } + + @Test + @DisplayName("Entity total should not decrease when already higher than success + failed") + void testEntityTotalNotDecreasedWhenAlreadyHigher() { + Set entities = Set.of("table"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn(Map.of("table", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + // Manually set a higher initial total to simulate real DB count + stats.getEntityStats().getAdditionalProperties().get("table").setTotalRecords(200); + stats.getJobStats().setTotalRecords(200); + stats.getReaderStats().setTotalRecords(200); + + executor.updateStats("table", new StepStats().withSuccessRecords(50).withFailedRecords(2)); + + Stats finalStats = executor.getStats().get(); + StepStats tableStats = finalStats.getEntityStats().getAdditionalProperties().get("table"); + + assertEquals(50, tableStats.getSuccessRecords()); + assertEquals(2, tableStats.getFailedRecords()); + // Total should remain 200 since 52 < 200 + assertEquals(200, tableStats.getTotalRecords()); + } + + @Test + @DisplayName("Reader total should be adjusted when job total exceeds it") + void testReaderTotalAdjustedFromJobTotal() { + Set entities = Set.of("table", "dashboard"); + + lenient() + .when(searchRepository.getEntityIndexMap()) + .thenReturn( + Map.of("table", mock(IndexMapping.class), "dashboard", mock(IndexMapping.class))); + + Stats stats = executor.initializeTotalRecords(entities); + executor.getStats().set(stats); + + // Simulate processing that exceeds initial totals + executor.updateStats("table", new StepStats().withSuccessRecords(60).withFailedRecords(5)); + executor.updateStats( + "dashboard", new StepStats().withSuccessRecords(30).withFailedRecords(2)); + + Stats finalStats = executor.getStats().get(); + + // Reader total should have been bumped to match the adjusted job total + int expectedTotal = 65 + 32; // table (60+5) + dashboard (30+2) + assertEquals(expectedTotal, finalStats.getReaderStats().getTotalRecords()); + assertEquals(expectedTotal, finalStats.getJobStats().getTotalRecords()); + } + } + + private ElasticSearchBulkSink.CustomBulkProcessor getCustomBulkProcessor( + ElasticSearchBulkSink sink) throws Exception { + java.lang.reflect.Field field = ElasticSearchBulkSink.class.getDeclaredField("bulkProcessor"); + field.setAccessible(true); + return (ElasticSearchBulkSink.CustomBulkProcessor) field.get(sink); + } + + private boolean invokeShouldRetry( + ElasticSearchBulkSink.CustomBulkProcessor processor, int attemptNumber, String errorMessage) + throws Exception { + Method method = + ElasticSearchBulkSink.CustomBulkProcessor.class.getDeclaredMethod( + "shouldRetry", int.class, Throwable.class); + method.setAccessible(true); + Throwable error = + errorMessage == null ? new RuntimeException() : new RuntimeException(errorMessage); + return (boolean) method.invoke(processor, attemptNumber, error); + } + + private boolean invokeIsPayloadTooLargeError( + ElasticSearchBulkSink.CustomBulkProcessor processor, String errorMessage) throws Exception { + Method method = + ElasticSearchBulkSink.CustomBulkProcessor.class.getDeclaredMethod( + "isPayloadTooLargeError", Throwable.class); + method.setAccessible(true); + Throwable error = + errorMessage == null ? new RuntimeException() : new RuntimeException(errorMessage); + return (boolean) method.invoke(processor, error); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategyTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategyTest.java new file mode 100644 index 000000000000..eefbdf241250 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/SingleServerIndexingStrategyTest.java @@ -0,0 +1,75 @@ +package org.openmetadata.service.apps.bundles.searchIndex; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import org.mockito.MockedConstruction; +import org.mockito.Mockito; +import org.openmetadata.schema.system.Stats; +import org.openmetadata.service.jdbi3.CollectionDAO; +import org.openmetadata.service.search.SearchRepository; + +class SingleServerIndexingStrategyTest { + + @Test + void delegatesExecutorOperations() { + CollectionDAO collectionDAO = mock(CollectionDAO.class); + SearchRepository searchRepository = mock(SearchRepository.class); + ReindexingProgressListener listener = mock(ReindexingProgressListener.class); + ReindexingJobContext context = mock(ReindexingJobContext.class); + ReindexingConfiguration config = + ReindexingConfiguration.builder().entities(java.util.Set.of("table")).build(); + ExecutionResult result = + new ExecutionResult(ExecutionResult.Status.COMPLETED, 10, 9, 1, 100, 200, new Stats()); + Stats stats = new Stats(); + + try (MockedConstruction mocked = + Mockito.mockConstruction( + SearchIndexExecutor.class, + (executor, mockContext) -> { + when(executor.addListener(listener)).thenReturn(executor); + when(executor.execute(config, context)).thenReturn(result); + when(executor.getStats()).thenReturn(new AtomicReference<>(stats)); + when(executor.isStopped()).thenReturn(true); + })) { + SingleServerIndexingStrategy strategy = + new SingleServerIndexingStrategy(collectionDAO, searchRepository); + + strategy.addListener(listener); + assertSame(result, strategy.execute(config, context)); + assertEquals(Optional.of(stats), strategy.getStats()); + strategy.stop(); + assertTrue(strategy.isStopped()); + + SearchIndexExecutor executor = mocked.constructed().get(0); + verify(executor).addListener(listener); + verify(executor).execute(config, context); + verify(executor).getStats(); + verify(executor).stop(); + verify(executor).isStopped(); + } + } + + @Test + void getStatsHandlesMissingExecutorStats() { + try (MockedConstruction mocked = + Mockito.mockConstruction( + SearchIndexExecutor.class, + (executor, mockContext) -> + when(executor.getStats()).thenReturn(new AtomicReference<>()))) { + SingleServerIndexingStrategy strategy = + new SingleServerIndexingStrategy(mock(CollectionDAO.class), mock(SearchRepository.class)); + + assertEquals(Optional.empty(), strategy.getStats()); + assertFalse(strategy.isStopped()); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContextTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContextTest.java index 72291680b35c..f4fe343c73b8 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContextTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobContextTest.java @@ -1,6 +1,7 @@ package org.openmetadata.service.apps.bundles.searchIndex.distributed; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.Map; import java.util.UUID; @@ -22,6 +23,7 @@ void distributedContextExposesJobMetadataAndCustomSource() { "DistributedSearchIndex-" + jobId.toString().substring(0, 8), context.getJobName()); assertEquals(200L, context.getStartTime()); assertEquals(jobId, context.getAppId()); + assertTrue(context.isDistributed()); assertEquals("REDIS", context.getSource()); assertEquals(job, context.getJob()); assertEquals(Map.of("participants", 3), context.getDistributedMetadata()); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactoryTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactoryTest.java index 1cdb4e022427..bed5c3bb13e8 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactoryTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobNotifierFactoryTest.java @@ -6,6 +6,7 @@ import java.lang.reflect.Constructor; import org.junit.jupiter.api.Test; +import org.openmetadata.service.cache.CacheConfig; import org.openmetadata.service.jdbi3.CollectionDAO; class DistributedJobNotifierFactoryTest { @@ -13,11 +14,29 @@ class DistributedJobNotifierFactoryTest { private final CollectionDAO collectionDAO = mock(CollectionDAO.class); @Test - void createUsesPollingNotifier() { + void createUsesRedisNotifierWhenRedisConfigIsComplete() { + CacheConfig cacheConfig = new CacheConfig(); + cacheConfig.provider = CacheConfig.Provider.redis; + cacheConfig.redis.url = "redis://cache:6379"; + DistributedJobNotifier notifier = - DistributedJobNotifierFactory.create(collectionDAO, "server-1"); + DistributedJobNotifierFactory.create(cacheConfig, collectionDAO, "server-1"); + + assertInstanceOf(RedisJobNotifier.class, notifier); + } + + @Test + void createFallsBackToPollingWhenRedisConfigIsMissingOrInvalid() { + CacheConfig missingUrlConfig = new CacheConfig(); + missingUrlConfig.provider = CacheConfig.Provider.redis; + + DistributedJobNotifier missingUrlNotifier = + DistributedJobNotifierFactory.create(missingUrlConfig, collectionDAO, "server-1"); + DistributedJobNotifier nullConfigNotifier = + DistributedJobNotifierFactory.create(null, collectionDAO, "server-1"); - assertInstanceOf(PollingJobNotifier.class, notifier); + assertInstanceOf(PollingJobNotifier.class, missingUrlNotifier); + assertInstanceOf(PollingJobNotifier.class, nullConfigNotifier); } @Test diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipantTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipantTest.java index 0bc79ab2abde..6237567558e4 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipantTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipantTest.java @@ -64,6 +64,7 @@ import org.openmetadata.service.Entity; import org.openmetadata.service.apps.bundles.searchIndex.BulkSink; import org.openmetadata.service.apps.bundles.searchIndex.IndexingFailureRecorder; +import org.openmetadata.service.cache.CacheConfig; import org.openmetadata.service.jdbi3.AppRepository; import org.openmetadata.service.jdbi3.CollectionDAO; import org.openmetadata.service.search.SearchClusterMetrics; @@ -145,7 +146,9 @@ void tearDown() throws Exception { @Test void testStartAndStop() { - participant = new DistributedJobParticipant(collectionDAO, searchRepository, "test-server-1"); + participant = + new DistributedJobParticipant( + collectionDAO, searchRepository, "test-server-1", (CacheConfig) null); // Initially not participating assertFalse(participant.isParticipating()); @@ -163,7 +166,9 @@ void testStartAndStop() { @Test void testMultipleStartCallsAreIdempotent() { - participant = new DistributedJobParticipant(collectionDAO, searchRepository, "test-server-1"); + participant = + new DistributedJobParticipant( + collectionDAO, searchRepository, "test-server-1", (CacheConfig) null); participant.start(); participant.start(); // Second call should be no-op @@ -177,7 +182,9 @@ void testMultipleStartCallsAreIdempotent() { @Test void testMultipleStopCallsAreIdempotent() { - participant = new DistributedJobParticipant(collectionDAO, searchRepository, "test-server-1"); + participant = + new DistributedJobParticipant( + collectionDAO, searchRepository, "test-server-1", (CacheConfig) null); participant.start(); participant.stop(); @@ -200,7 +207,9 @@ void testDoesNotJoinWhenNoRunningJobs() throws Exception { DistributedSearchIndexCoordinator.class, (mock, context) -> when(mock.getRecentJobs(any(), anyInt())).thenReturn(List.of()))) { - participant = new DistributedJobParticipant(collectionDAO, searchRepository, "test-server-1"); + participant = + new DistributedJobParticipant( + collectionDAO, searchRepository, "test-server-1", (CacheConfig) null); participant.start(); // Wait a bit for the scheduler to run at least once @@ -228,7 +237,6 @@ void testJoinsActiveJobWithPendingPartitions() { .id(jobId) .status(IndexJobStatus.RUNNING) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .totalRecords(100) .build(); @@ -237,7 +245,6 @@ void testJoinsActiveJobWithPendingPartitions() { .id(jobId) .status(IndexJobStatus.COMPLETED) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .totalRecords(100) .processedRecords(100) .successRecords(100) @@ -322,7 +329,6 @@ void testDoesNotRejoinSameRunningJob() { .id(jobId) .status(IndexJobStatus.RUNNING) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .totalRecords(100) .build(); @@ -395,7 +401,6 @@ void testClearsJobIdWhenJobCompletes() { .id(jobId) .status(IndexJobStatus.RUNNING) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .totalRecords(100) .build(); @@ -404,7 +409,6 @@ void testClearsJobIdWhenJobCompletes() { .id(jobId) .status(IndexJobStatus.COMPLETED) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .totalRecords(100) .processedRecords(100) .successRecords(100) @@ -494,7 +498,6 @@ void testAttemptsToClaimPartitions() { .id(jobId) .status(IndexJobStatus.RUNNING) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .totalRecords(100) .build(); @@ -503,7 +506,6 @@ void testAttemptsToClaimPartitions() { .id(jobId) .status(IndexJobStatus.COMPLETED) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .totalRecords(100) .processedRecords(100) .successRecords(100) @@ -666,7 +668,6 @@ void testJoinAndProcessJobTracksPollingNotifierParticipation() throws Exception .id(jobId) .status(IndexJobStatus.RUNNING) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .build(); SearchIndexPartition pendingPartition = @@ -846,6 +847,7 @@ void testRecoveredParticipationRestoresRunRecordAndFinalizesStats() throws Excep config.setBatchSize(50); config.setMaxConcurrentRequests(8); config.setPayLoadSize(4096L); + config.setRecreateIndex(true); SearchIndexJob runningJob = SearchIndexJob.builder() @@ -880,7 +882,7 @@ void testRecoveredParticipationRestoresRunRecordAndFinalizesStats() throws Excep CollectionDAO.AppExtensionTimeSeries appExtensionDao = mock(CollectionDAO.AppExtensionTimeSeries.class); AtomicReference callbackRef = new AtomicReference<>(); - AtomicReference stagedIndexContextRef = new AtomicReference<>(); + AtomicReference recreateContextRef = new AtomicReference<>(); SuccessContext successContext = new SuccessContext().withAdditionalProperty("recovered", "yes"); when(appRepository.getDao()).thenReturn(appDao); @@ -926,7 +928,7 @@ void testRecoveredParticipationRestoresRunRecordAndFinalizesStats() throws Excep mockConstruction( PartitionWorker.class, (mock, context) -> { - stagedIndexContextRef.set(context.arguments().get(3)); + recreateContextRef.set(context.arguments().get(3)); when(mock.processPartition(partition)) .thenReturn(new PartitionWorker.PartitionResult(4, 1, false, 2, 3)); }); @@ -945,7 +947,7 @@ void testRecoveredParticipationRestoresRunRecordAndFinalizesStats() throws Excep "processJobPartitions", new Class[] {SearchIndexJob.class}, runningJob); assertNotNull(callbackRef.get()); - assertNotNull(stagedIndexContextRef.get()); + assertNotNull(recreateContextRef.get()); callbackRef .get() .onFailure( @@ -1003,38 +1005,6 @@ void testRecoveredParticipationRestoresRunRecordAndFinalizesStats() throws Excep } } - @Test - void testProcessJobPartitionsSkipsJobWithoutStagedIndexMapping() throws Exception { - UUID jobId = UUID.randomUUID(); - EventPublisherJob config = new EventPublisherJob(); - config.setEntities(Set.of("table")); - - SearchIndexJob runningJob = - SearchIndexJob.builder() - .id(jobId) - .status(IndexJobStatus.RUNNING) - .jobConfiguration(config) - .build(); - - participant = - new DistributedJobParticipant( - collectionDAO, searchRepository, "test-server-1", testNotifier); - setParticipantRunning(true); - - try (MockedConstruction failureConstruction = - mockConstruction(IndexingFailureRecorder.class); - MockedConstruction workerConstruction = - mockConstruction(PartitionWorker.class)) { - - invokeParticipantMethod( - "processJobPartitions", new Class[] {SearchIndexJob.class}, runningJob); - - verify(searchRepository, never()).createBulkSink(anyInt(), anyInt(), anyLong()); - assertTrue(failureConstruction.constructed().isEmpty()); - assertTrue(workerConstruction.constructed().isEmpty()); - } - } - @Test void testProcessJobPartitionsUsesDefaultBulkSinkSettingsAndHandlesInterruptedWait() throws Exception { @@ -1047,7 +1017,6 @@ void testProcessJobPartitionsUsesDefaultBulkSinkSettingsAndHandlesInterruptedWai .id(jobId) .status(IndexJobStatus.RUNNING) .jobConfiguration(config) - .stagedIndexMapping(Map.of("table", "table_staged")) .build(); SearchIndexPartition pendingPartition = SearchIndexPartition.builder() diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutorTest.java index e3526f2da4fa..e2252ab850fd 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutorTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutorTest.java @@ -372,7 +372,7 @@ void getFreshStatsAndUpdateStagedIndexMappingUseCurrentJob() throws Exception { @Test void initializeEntityTrackerCountsPartitionsAndWiresPromotionCallback() throws Exception { UUID jobId = UUID.randomUUID(); - ReindexContext stagedIndexContext = mock(ReindexContext.class); + ReindexContext recreateContext = mock(ReindexContext.class); SearchRepository searchRepository = mock(SearchRepository.class); RecreateIndexHandler recreateHandler = mock(RecreateIndexHandler.class); @@ -382,48 +382,50 @@ void initializeEntityTrackerCountsPartitionsAndWiresPromotionCallback() throws E partition(jobId, "table", PartitionStatus.PENDING), partition(jobId, "table", PartitionStatus.COMPLETED), partition(jobId, "dashboard", PartitionStatus.FAILED))); - when(stagedIndexContext.getEntities()).thenReturn(Set.of("table", "dashboard")); + when(recreateContext.getEntities()).thenReturn(Set.of("table", "dashboard")); setField("entityTracker", new EntityCompletionTracker(jobId)); - setField("stagedIndexContext", stagedIndexContext); + setField("recreateContext", recreateContext); try (MockedStatic entityMock = mockStatic(Entity.class)) { entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); when(searchRepository.createReindexHandler()).thenReturn(recreateHandler); - invokePrivate("initializeEntityTracker", new Class[] {UUID.class}, jobId); + invokePrivate( + "initializeEntityTracker", new Class[] {UUID.class, boolean.class}, jobId, true); } EntityCompletionTracker tracker = executor.getEntityTracker(); assertNotNull(tracker); assertEquals(2, tracker.getStatus("table").totalPartitions()); assertEquals(1, tracker.getStatus("dashboard").totalPartitions()); - assertSame(recreateHandler, getField("indexPromotionHandler")); + assertSame(recreateHandler, getField("recreateIndexHandler")); } @Test void initializeEntityTrackerCallbackPromotesEntityWhenTrackingCompletes() throws Exception { UUID jobId = UUID.randomUUID(); - ReindexContext stagedIndexContext = mock(ReindexContext.class); + ReindexContext recreateContext = mock(ReindexContext.class); DefaultRecreateHandler recreateHandler = mock(DefaultRecreateHandler.class); SearchRepository searchRepository = mock(SearchRepository.class); when(coordinator.getPartitions(jobId, null)) .thenReturn(List.of(partition(jobId, "table", PartitionStatus.PENDING))); - when(stagedIndexContext.getEntities()).thenReturn(Set.of("table")); - when(stagedIndexContext.getStagedIndex("table")).thenReturn(Optional.of("staged_table")); - when(stagedIndexContext.getCanonicalIndex("table")).thenReturn(Optional.of("table_search")); - when(stagedIndexContext.getOriginalIndex("table")).thenReturn(Optional.of("table_current")); - when(stagedIndexContext.getCanonicalAlias("table")).thenReturn(Optional.of("table_alias")); - when(stagedIndexContext.getExistingAliases("table")).thenReturn(Set.of("table_existing")); - when(stagedIndexContext.getParentAliases("table")).thenReturn(List.of("table_parent")); + when(recreateContext.getEntities()).thenReturn(Set.of("table")); + when(recreateContext.getStagedIndex("table")).thenReturn(Optional.of("staged_table")); + when(recreateContext.getCanonicalIndex("table")).thenReturn(Optional.of("table_search")); + when(recreateContext.getOriginalIndex("table")).thenReturn(Optional.of("table_current")); + when(recreateContext.getCanonicalAlias("table")).thenReturn(Optional.of("table_alias")); + when(recreateContext.getExistingAliases("table")).thenReturn(Set.of("table_existing")); + when(recreateContext.getParentAliases("table")).thenReturn(List.of("table_parent")); setField("entityTracker", new EntityCompletionTracker(jobId)); - setField("stagedIndexContext", stagedIndexContext); + setField("recreateContext", recreateContext); try (MockedStatic entityMock = mockStatic(Entity.class)) { entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); when(searchRepository.createReindexHandler()).thenReturn(recreateHandler); - invokePrivate("initializeEntityTracker", new Class[] {UUID.class}, jobId); + invokePrivate( + "initializeEntityTracker", new Class[] {UUID.class, boolean.class}, jobId, true); } executor.getEntityTracker().recordPartitionComplete("table", false); @@ -433,18 +435,18 @@ void initializeEntityTrackerCallbackPromotesEntityWhenTrackingCompletes() throws @Test void promoteEntityIndexUsesDefaultAndGenericHandlers() throws Exception { - ReindexContext stagedIndexContext = mock(ReindexContext.class); + ReindexContext recreateContext = mock(ReindexContext.class); DefaultRecreateHandler defaultHandler = mock(DefaultRecreateHandler.class); RecreateIndexHandler genericHandler = mock(RecreateIndexHandler.class); - when(stagedIndexContext.getStagedIndex("table")).thenReturn(Optional.of("staged_table")); - when(stagedIndexContext.getCanonicalIndex("table")).thenReturn(Optional.of("table_search")); - when(stagedIndexContext.getOriginalIndex("table")).thenReturn(Optional.of("table_current")); - when(stagedIndexContext.getCanonicalAlias("table")).thenReturn(Optional.of("table_alias")); - when(stagedIndexContext.getExistingAliases("table")).thenReturn(Set.of("table_existing")); - when(stagedIndexContext.getParentAliases("table")).thenReturn(List.of("table_parent")); + when(recreateContext.getStagedIndex("table")).thenReturn(Optional.of("staged_table")); + when(recreateContext.getCanonicalIndex("table")).thenReturn(Optional.of("table_search")); + when(recreateContext.getOriginalIndex("table")).thenReturn(Optional.of("table_current")); + when(recreateContext.getCanonicalAlias("table")).thenReturn(Optional.of("table_alias")); + when(recreateContext.getExistingAliases("table")).thenReturn(Set.of("table_existing")); + when(recreateContext.getParentAliases("table")).thenReturn(List.of("table_parent")); - setField("stagedIndexContext", stagedIndexContext); - setField("indexPromotionHandler", defaultHandler); + setField("recreateContext", recreateContext); + setField("recreateIndexHandler", defaultHandler); invokePrivate( "promoteEntityIndex", new Class[] {String.class, boolean.class}, "table", false); @@ -456,12 +458,12 @@ void promoteEntityIndexUsesDefaultAndGenericHandlers() throws Exception { assertEquals("staged_table", contextCaptor.getValue().getStagedIndex()); assertTrue(contextCaptor.getValue().getParentAliases().contains("table_parent")); - setField("indexPromotionHandler", genericHandler); + setField("recreateIndexHandler", genericHandler); invokePrivate( "promoteEntityIndex", new Class[] {String.class, boolean.class}, "table", true); verify(genericHandler).finalizeReindex(any(EntityReindexContext.class), eq(true)); - when(stagedIndexContext.getStagedIndex("topic")).thenReturn(Optional.empty()); + when(recreateContext.getStagedIndex("topic")).thenReturn(Optional.empty()); invokePrivate( "promoteEntityIndex", new Class[] {String.class, boolean.class}, "topic", true); verifyNoMoreInteractions(genericHandler); @@ -472,20 +474,20 @@ void promoteEntityIndexReturnsWithoutContextAndSwallowsHandlerFailures() throws invokePrivate( "promoteEntityIndex", new Class[] {String.class, boolean.class}, "table", true); - ReindexContext stagedIndexContext = mock(ReindexContext.class); + ReindexContext recreateContext = mock(ReindexContext.class); DefaultRecreateHandler defaultHandler = mock(DefaultRecreateHandler.class); - when(stagedIndexContext.getStagedIndex("table")).thenReturn(Optional.of("staged_table")); - when(stagedIndexContext.getCanonicalIndex("table")).thenReturn(Optional.of("table_search")); - when(stagedIndexContext.getOriginalIndex("table")).thenReturn(Optional.of("table_current")); - when(stagedIndexContext.getCanonicalAlias("table")).thenReturn(Optional.of("table_alias")); - when(stagedIndexContext.getExistingAliases("table")).thenReturn(Set.of()); - when(stagedIndexContext.getParentAliases("table")).thenReturn(List.of()); + when(recreateContext.getStagedIndex("table")).thenReturn(Optional.of("staged_table")); + when(recreateContext.getCanonicalIndex("table")).thenReturn(Optional.of("table_search")); + when(recreateContext.getOriginalIndex("table")).thenReturn(Optional.of("table_current")); + when(recreateContext.getCanonicalAlias("table")).thenReturn(Optional.of("table_alias")); + when(recreateContext.getExistingAliases("table")).thenReturn(Set.of()); + when(recreateContext.getParentAliases("table")).thenReturn(List.of()); doThrow(new IllegalStateException("promotion failed")) .when(defaultHandler) .promoteEntityIndex(any(EntityReindexContext.class), eq(true)); - setField("stagedIndexContext", stagedIndexContext); - setField("indexPromotionHandler", defaultHandler); + setField("recreateContext", recreateContext); + setField("recreateIndexHandler", defaultHandler); invokePrivate( "promoteEntityIndex", new Class[] {String.class, boolean.class}, "table", true); @@ -571,7 +573,8 @@ void executeDoesNotNotifyPeersWhenStartedJobIsNotRunning() throws Exception { () -> executor.execute( bulkSink, - stagedContext("table"), + null, + false, ReindexingConfiguration.builder().entities(Set.of("table")).build())); assertTrue(exception.getMessage().contains(IndexJobStatus.FAILED.name())); @@ -616,7 +619,8 @@ void executeDoesNotRebroadcastStartWhenJoiningRunningJob() throws Exception { DistributedSearchIndexExecutor.ExecutionResult result = executor.execute( bulkSink, - stagedContext("table"), + null, + false, ReindexingConfiguration.builder() .entities(Set.of("table")) .consumerThreads(1) @@ -636,7 +640,7 @@ void executeRequiresCurrentJobBeforeRunning() { IllegalStateException.class, () -> executor.execute( - mock(BulkSink.class), null, ReindexingConfiguration.builder().build())); + mock(BulkSink.class), null, false, ReindexingConfiguration.builder().build())); assertTrue(exception.getMessage().contains("No job to execute")); } @@ -701,7 +705,8 @@ void executeRunsMinimalHappyPathAndCleansUpResources() throws Exception { DistributedSearchIndexExecutor.ExecutionResult result = executor.execute( bulkSink, - stagedContext("table"), + null, + false, ReindexingConfiguration.builder() .entities(Set.of("table")) .consumerThreads(1) @@ -749,9 +754,7 @@ void executeRecordsFailuresFromCleanupWithoutAbortingResult() throws Exception { runningJob.withStatus(IndexJobStatus.FAILED).withFailedRecords(2).withCompletedAt(400L); BulkSink bulkSink = mock(BulkSink.class); ReindexingProgressListener listener = mock(ReindexingProgressListener.class); - ReindexContext stagedIndexContext = mock(ReindexContext.class); - SearchRepository searchRepository = mock(SearchRepository.class); - RecreateIndexHandler indexPromotionHandler = mock(RecreateIndexHandler.class); + ReindexContext recreateContext = mock(ReindexContext.class); ReindexingMetrics metrics = mock(ReindexingMetrics.class); Timer.Sample timerSample = mock(Timer.Sample.class); AtomicReference callbackRef = new AtomicReference<>(); @@ -796,18 +799,16 @@ void executeRecordsFailuresFromCleanupWithoutAbortingResult() throws Exception { IndexingFailureRecorder.class, (mock, context) -> doThrow(new IllegalStateException("close failed")).when(mock).close()); - MockedStatic entityMock = mockStatic(Entity.class); MockedStatic metricsMock = mockStatic(ReindexingMetrics.class)) { - entityMock.when(Entity::getSearchRepository).thenReturn(searchRepository); - when(searchRepository.createReindexHandler()).thenReturn(indexPromotionHandler); metricsMock.when(ReindexingMetrics::getInstance).thenReturn(metrics); when(metrics.startJobTimer()).thenReturn(timerSample); DistributedSearchIndexExecutor.ExecutionResult result = executor.execute( bulkSink, - stagedIndexContext, + recreateContext, + false, ReindexingConfiguration.builder() .entities(Set.of("table")) .consumerThreads(1) @@ -885,7 +886,8 @@ void executeHandlesInterruptedAwaitAndStoppedMetricsCleanup() throws Exception { DistributedSearchIndexExecutor.ExecutionResult result = executor.execute( bulkSink, - stagedContext("table"), + null, + false, ReindexingConfiguration.builder() .entities(Set.of("table")) .consumerThreads(1) @@ -954,7 +956,8 @@ void executeForceCompletesStoppingJobsDuringCleanupAndRecordsStoppedMetrics() th DistributedSearchIndexExecutor.ExecutionResult result = executor.execute( bulkSink, - stagedContext("table"), + null, + false, ReindexingConfiguration.builder() .entities(Set.of("table")) .consumerThreads(1) @@ -998,6 +1001,7 @@ void runWorkerLoopAggregatesPartitionResults() throws Exception { BulkSink.class, int.class, ReindexContext.class, + boolean.class, AtomicLong.class, AtomicLong.class, ReindexingConfiguration.class @@ -1005,7 +1009,8 @@ void runWorkerLoopAggregatesPartitionResults() throws Exception { 0, bulkSink, 100, - stagedContext("table"), + null, + false, totalSuccess, totalFailed, ReindexingConfiguration.builder().build()); @@ -1052,6 +1057,7 @@ void runWorkerLoopRetriesClaimingAndBreaksOnInterruptedSleep() throws Exception BulkSink.class, int.class, ReindexContext.class, + boolean.class, AtomicLong.class, AtomicLong.class, ReindexingConfiguration.class @@ -1059,7 +1065,8 @@ void runWorkerLoopRetriesClaimingAndBreaksOnInterruptedSleep() throws Exception 2, mock(BulkSink.class), 100, - stagedContext("table"), + null, + false, new AtomicLong(), new AtomicLong(), ReindexingConfiguration.builder().build()); @@ -1104,6 +1111,7 @@ void runWorkerLoopSwallowsPartitionProcessingErrorsAndCleansUpState() throws Exc BulkSink.class, int.class, ReindexContext.class, + boolean.class, AtomicLong.class, AtomicLong.class, ReindexingConfiguration.class @@ -1111,7 +1119,8 @@ void runWorkerLoopSwallowsPartitionProcessingErrorsAndCleansUpState() throws Exc 1, mock(BulkSink.class), 100, - stagedContext("table"), + null, + false, new AtomicLong(), new AtomicLong(), ReindexingConfiguration.builder().build()); @@ -1261,19 +1270,6 @@ private SearchIndexPartition partition(UUID jobId, String entityType, PartitionS .build(); } - private ReindexContext stagedContext(String entityType) { - ReindexContext context = new ReindexContext(); - context.add( - entityType, - entityType + "_index", - entityType + "_original", - entityType + "_staged", - Set.of(), - entityType, - List.of()); - return context; - } - private Object invokePrivate(String methodName, Class[] parameterTypes, Object... args) throws Exception { Method method = diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java index 537dfe8b0ee1..d554f33fc4d6 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PartitionWorkerTest.java @@ -64,7 +64,6 @@ import org.openmetadata.service.apps.bundles.searchIndex.BulkSink; import org.openmetadata.service.apps.bundles.searchIndex.IndexingFailureRecorder; import org.openmetadata.service.apps.bundles.searchIndex.ReindexingConfiguration; -import org.openmetadata.service.apps.bundles.searchIndex.SearchIndexEntityTypes; import org.openmetadata.service.apps.bundles.searchIndex.stats.StageCounter; import org.openmetadata.service.apps.bundles.searchIndex.stats.StageStatsTracker; import org.openmetadata.service.exception.SearchIndexException; @@ -84,7 +83,7 @@ class PartitionWorkerTest { @Mock private CollectionDAO collectionDAO; @Mock private CollectionDAO.SearchIndexServerStatsDAO searchIndexServerStatsDAO; @Mock private BulkSink bulkSink; - @Mock private ReindexContext stagedIndexContext; + @Mock private ReindexContext recreateContext; @Mock private ReindexingConfiguration reindexingConfiguration; private PartitionWorker worker; @@ -94,10 +93,7 @@ class PartitionWorkerTest { @BeforeEach void setUp() { - when(stagedIndexContext.getStagedIndex(any())) - .thenAnswer( - invocation -> Optional.of(invocation.getArgument(0, String.class) + "_staging")); - worker = new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext); + worker = new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, recreateContext, false); } @Test @@ -148,14 +144,14 @@ void testPartitionResult_WasStopped() { @Test void testWorkerWithDifferentConfigurations() { PartitionWorker workerWithRecreate = - new PartitionWorker(coordinator, bulkSink, 200, stagedIndexContext); + new PartitionWorker(coordinator, bulkSink, 200, recreateContext, true); assertFalse(workerWithRecreate.isStopped()); - PartitionWorker workerWithSmallBatch = - new PartitionWorker(coordinator, bulkSink, 50, stagedIndexContext); + PartitionWorker workerWithoutContext = + new PartitionWorker(coordinator, bulkSink, 50, null, false); - assertFalse(workerWithSmallBatch.isStopped()); + assertFalse(workerWithoutContext.isStopped()); } @Test @@ -401,17 +397,17 @@ void initializeKeysetCursorHitsPrecomputedCacheAndSkipsOffsetFallback() throws E } @Test - void createContextDataIncludesStagedContextTargetIndexAndStatsTracker() throws Exception { - PartitionWorker stagedWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext); + void createContextDataIncludesRecreateContextTargetIndexAndStatsTracker() throws Exception { + PartitionWorker recreateWorker = + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, recreateContext, true); StageStatsTracker statsTracker = mock(StageStatsTracker.class); - when(stagedIndexContext.getStagedIndex("table")).thenReturn(Optional.of("table_staging")); + when(recreateContext.getStagedIndex("table")).thenReturn(Optional.of("table_staging")); @SuppressWarnings("unchecked") Map contextData = (Map) invokePrivate( - stagedWorker, + recreateWorker, "createContextData", new Class[] {String.class, StageStatsTracker.class}, "table", @@ -420,37 +416,16 @@ void createContextDataIncludesStagedContextTargetIndexAndStatsTracker() throws E assertEquals("table", contextData.get("entityType")); assertEquals(Boolean.TRUE, contextData.get("recreateIndex")); assertEquals(statsTracker, contextData.get(BulkSink.STATS_TRACKER_CONTEXT_KEY)); - assertEquals(stagedIndexContext, contextData.get("recreateContext")); + assertEquals(recreateContext, contextData.get("recreateContext")); assertEquals("table_staging", contextData.get("targetIndex")); } - @Test - void createContextDataNormalizesLegacyEntityAliasesBeforeStagedIndexLookup() throws Exception { - when(stagedIndexContext.getStagedIndex(Entity.QUERY_COST_RECORD)) - .thenReturn(Optional.of("query_cost_record_staging")); - - @SuppressWarnings("unchecked") - Map contextData = - (Map) - invokePrivate( - worker, - "createContextData", - new Class[] {String.class, StageStatsTracker.class}, - SearchIndexEntityTypes.QUERY_COST_RESULT, - null); - - assertEquals(Entity.QUERY_COST_RECORD, contextData.get("entityType")); - assertEquals("query_cost_record_staging", contextData.get("targetIndex")); - verify(stagedIndexContext).getStagedIndex(Entity.QUERY_COST_RECORD); - verify(stagedIndexContext, never()).getStagedIndex(SearchIndexEntityTypes.QUERY_COST_RESULT); - } - @Test void processBatchWritesEntitiesAndRecordsReaderFailures() throws Exception { IndexingFailureRecorder failureRecorder = mock(IndexingFailureRecorder.class); StageStatsTracker statsTracker = mock(StageStatsTracker.class); PartitionWorker batchWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext, failureRecorder); + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, null, false, failureRecorder); EntityInterface entityOne = mock(EntityInterface.class); EntityInterface entityTwo = mock(EntityInterface.class); @@ -489,7 +464,7 @@ void processBatchWritesEntitiesAndRecordsReaderFailures() throws Exception { verify(bulkSink).write(entitiesCaptor.capture(), contextCaptor.capture()); assertEquals(List.of(entityOne, entityTwo), entitiesCaptor.getValue()); assertEquals("table", contextCaptor.getValue().get("entityType")); - assertEquals(Boolean.TRUE, contextCaptor.getValue().get("recreateIndex")); + assertEquals(Boolean.FALSE, contextCaptor.getValue().get("recreateIndex")); assertEquals(statsTracker, contextCaptor.getValue().get(BulkSink.STATS_TRACKER_CONTEXT_KEY)); } @@ -498,7 +473,7 @@ void processBatchExtractsIdFromEntityInterfaceForReaderFailure() throws Exceptio IndexingFailureRecorder failureRecorder = mock(IndexingFailureRecorder.class); StageStatsTracker statsTracker = mock(StageStatsTracker.class); PartitionWorker batchWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext, failureRecorder); + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, null, false, failureRecorder); UUID errorEntityId = UUID.randomUUID(); EntityInterface failingEntity = mock(EntityInterface.class); @@ -529,7 +504,7 @@ void processBatchSkipsReaderFailureWhenEntityInterfaceHasNullId() throws Excepti IndexingFailureRecorder failureRecorder = mock(IndexingFailureRecorder.class); StageStatsTracker statsTracker = mock(StageStatsTracker.class); PartitionWorker batchWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext, failureRecorder); + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, null, false, failureRecorder); EntityInterface failingEntity = mock(EntityInterface.class); when(failingEntity.getId()).thenReturn(null); @@ -556,7 +531,7 @@ void processBatchSkipsReaderFailureWhenEntityInterfaceHasNullId() throws Excepti @Test void processBatchWrapsSinkFailuresAsSearchIndexException() throws Exception { PartitionWorker batchWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext); + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, null, false); ResultList resultList = new ResultList<>(); resultList.setData(List.of(mock(EntityInterface.class))); @@ -642,7 +617,7 @@ void readEntitiesKeysetPassesSelectiveFieldsNotWildcard() throws Exception { void readEntitiesKeysetUsesTimeSeriesSourceWithConfiguredWindow() throws Exception { PartitionWorker timeSeriesWorker = new PartitionWorker( - coordinator, bulkSink, BATCH_SIZE, stagedIndexContext, null, reindexingConfiguration); + coordinator, bulkSink, BATCH_SIZE, null, false, null, reindexingConfiguration); when(reindexingConfiguration.getTimeSeriesStartTs(Entity.QUERY_COST_RECORD)).thenReturn(100L); ResultList resultList = new ResultList<>(); @@ -675,43 +650,6 @@ void readEntitiesKeysetUsesTimeSeriesSourceWithConfiguredWindow() throws Excepti assertNotNull(constructorArgs.get().get(4)); } - @Test - void readEntitiesKeysetNormalizesLegacyTimeSeriesAliases() throws Exception { - PartitionWorker timeSeriesWorker = - new PartitionWorker( - coordinator, bulkSink, BATCH_SIZE, stagedIndexContext, null, reindexingConfiguration); - when(reindexingConfiguration.getTimeSeriesStartTs(Entity.QUERY_COST_RECORD)).thenReturn(100L); - - ResultList resultList = new ResultList<>(); - resultList.setData(List.of(mock(EntityTimeSeriesInterface.class))); - AtomicReference> constructorArgs = new AtomicReference<>(); - - try (MockedConstruction ignored = - mockConstruction( - PaginatedEntityTimeSeriesSource.class, - (mock, context) -> { - constructorArgs.set(List.copyOf(context.arguments())); - doReturn(resultList).when(mock).readWithCursor("cursor"); - })) { - - assertEquals( - resultList, - invokePrivate( - timeSeriesWorker, - "readEntitiesKeyset", - new Class[] {String.class, String.class, int.class}, - SearchIndexEntityTypes.QUERY_COST_RESULT, - "cursor", - 3)); - } - - assertEquals(Entity.QUERY_COST_RECORD, constructorArgs.get().get(0)); - assertEquals(3, constructorArgs.get().get(1)); - assertEquals(List.of(), constructorArgs.get().get(2)); - assertEquals(100L, constructorArgs.get().get(3)); - assertNotNull(constructorArgs.get().get(4)); - } - @Test void writeToSinkUsesTimeSeriesEntitiesForTimeSeriesTypes() throws Exception { ResultList resultList = new ResultList<>(); @@ -753,7 +691,7 @@ void waitForSinkOperationsReconcilesStalePendingWorkAndFlushesStats() throws Exc @Test void processPartitionKeepsProgressStatusProcessingAndCompletesSuccessfully() { PartitionWorker partitionWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext); + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, null, false); SearchIndexPartition partition = buildPartition("table", 0, 2); ResultList resultList = new ResultList<>(); @@ -795,7 +733,7 @@ void processPartitionKeepsProgressStatusProcessingAndCompletesSuccessfully() { @Test void processPartitionTracksReaderFailuresAndCompletesWithFailedCounts() { PartitionWorker partitionWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext); + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, null, false); SearchIndexPartition partition = buildPartition("table", 0, 2); SearchIndexException readerFailure = @@ -835,7 +773,7 @@ void processPartitionTracksReaderFailuresAndCompletesWithFailedCounts() { @Test void processPartitionStopsAfterReadWhenStopRequestedMidLoop() { PartitionWorker partitionWorker = - new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, stagedIndexContext); + new PartitionWorker(coordinator, bulkSink, BATCH_SIZE, null, false); SearchIndexPartition partition = buildPartition("table", 0, 2); ResultList resultList = new ResultList<>(); @@ -879,7 +817,7 @@ void processPartitionStopsAfterReadWhenStopRequestedMidLoop() { void processPartitionRecordsSinkFailuresAndStopsWhenCursorCannotBeRebuilt() throws Exception { IndexingFailureRecorder failureRecorder = mock(IndexingFailureRecorder.class); PartitionWorker partitionWorker = - new PartitionWorker(coordinator, bulkSink, 2, stagedIndexContext, failureRecorder); + new PartitionWorker(coordinator, bulkSink, 2, null, false, failureRecorder); SearchIndexPartition partition = buildPartition("table", 0, 4); ResultList resultList = new ResultList<>(); @@ -929,8 +867,7 @@ void processPartitionRecordsSinkFailuresAndStopsWhenCursorCannotBeRebuilt() thro @Test void processPartitionAdjustsSuccessCountsForProcessFailures() { - PartitionWorker partitionWorker = - new PartitionWorker(coordinator, bulkSink, 2, stagedIndexContext); + PartitionWorker partitionWorker = new PartitionWorker(coordinator, bulkSink, 2, null, false); SearchIndexPartition partition = buildPartition("table", 0, 2); ResultList resultList = new ResultList<>(); @@ -974,8 +911,7 @@ void processPartitionAdjustsSuccessCountsForProcessFailures() { @Test void processPartitionFailsPartitionWhenCompletionThrows() { - PartitionWorker partitionWorker = - new PartitionWorker(coordinator, bulkSink, 2, stagedIndexContext); + PartitionWorker partitionWorker = new PartitionWorker(coordinator, bulkSink, 2, null, false); SearchIndexPartition partition = buildPartition("table", 0, 1); ResultList resultList = new ResultList<>(); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifierTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifierTest.java index 97405b7e6985..030438f487fa 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifierTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/PollingJobNotifierTest.java @@ -5,7 +5,6 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; @@ -68,14 +67,13 @@ void pollForJobsDiscoversNewJobsAndClearsCompletedOnes() throws Exception { AtomicReference callbackJob = new AtomicReference<>(); notifier.onJobStarted(callbackJob::set); getRunningFlag(notifier).set(true); - setFastIdleUntil(notifier, System.currentTimeMillis() + 60_000L); invokePoll(notifier); assertEquals(jobId, callbackJob.get()); assertTrue(getKnownJobs(notifier).contains(jobId)); - setLastPollTime(notifier, System.currentTimeMillis() - 5_000L); + setLastPollTime(notifier, System.currentTimeMillis() - 31_000L); invokePoll(notifier); assertTrue(getKnownJobs(notifier).isEmpty()); @@ -102,29 +100,6 @@ void activePollingIntervalAndExceptionHandlingBehaveAsExpected() throws Exceptio assertTrue(getKnownJobs(notifier).isEmpty()); } - @Test - void idlePollingBacksOffAfterFastWindowAndResumesAfterJobActivity() throws Exception { - CollectionDAO collectionDAO = mock(CollectionDAO.class); - CollectionDAO.SearchIndexJobDAO jobDAO = mock(CollectionDAO.SearchIndexJobDAO.class); - when(collectionDAO.searchIndexJobDAO()).thenReturn(jobDAO); - when(jobDAO.getRunningJobIds()).thenReturn(List.of()); - - PollingJobNotifier notifier = new PollingJobNotifier(collectionDAO, "server-backoff"); - getRunningFlag(notifier).set(true); - - setFastIdleUntil(notifier, System.currentTimeMillis() - 1L); - setLastPollTime(notifier, System.currentTimeMillis() - 10_000L); - invokePoll(notifier); - - verifyNoInteractions(jobDAO); - - notifier.notifyJobCompleted(UUID.randomUUID()); - setLastPollTime(notifier, System.currentTimeMillis() - 5_000L); - invokePoll(notifier); - - verify(jobDAO).getRunningJobIds(); - } - @Test void stopHandlesPreconfiguredSchedulerAndInterruptedTermination() throws Exception { PollingJobNotifier notifier = @@ -162,12 +137,6 @@ private void setLastPollTime(PollingJobNotifier notifier, long value) throws Exc field.setLong(notifier, value); } - private void setFastIdleUntil(PollingJobNotifier notifier, long value) throws Exception { - Field field = notifier.getClass().getDeclaredField("fastIdleUntil"); - field.setAccessible(true); - field.setLong(notifier, value); - } - private Object getField(Object target, String name) throws Exception { Field field = target.getClass().getDeclaredField(name); field.setAccessible(true); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/RedisJobNotifierTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/RedisJobNotifierTest.java new file mode 100644 index 000000000000..5eeab373a73f --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/RedisJobNotifierTest.java @@ -0,0 +1,300 @@ +package org.openmetadata.service.apps.bundles.searchIndex.distributed; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +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.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.lettuce.core.RedisClient; +import io.lettuce.core.RedisURI; +import io.lettuce.core.api.StatefulRedisConnection; +import io.lettuce.core.api.sync.RedisCommands; +import io.lettuce.core.pubsub.RedisPubSubAdapter; +import io.lettuce.core.pubsub.RedisPubSubListener; +import io.lettuce.core.pubsub.StatefulRedisPubSubConnection; +import io.lettuce.core.pubsub.api.sync.RedisPubSubCommands; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.time.Duration; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import org.openmetadata.service.cache.CacheConfig; + +class RedisJobNotifierTest { + + @Test + void startInitializesRedisConnectionsAndStopCleansThemUp() { + CacheConfig config = cacheConfig("redis://cache:6380"); + RedisJobNotifier notifier = new RedisJobNotifier(config, "server-123"); + RedisClient redisClient = mock(RedisClient.class); + StatefulRedisPubSubConnection subConnection = + mock(StatefulRedisPubSubConnection.class); + StatefulRedisConnection pubConnection = mock(StatefulRedisConnection.class); + RedisPubSubCommands pubSubCommands = mock(RedisPubSubCommands.class); + RedisCommands redisCommands = mock(RedisCommands.class); + when(redisClient.connectPubSub()).thenReturn(subConnection); + when(redisClient.connect()).thenReturn(pubConnection); + when(subConnection.sync()).thenReturn(pubSubCommands); + when(pubConnection.sync()).thenReturn(redisCommands); + + try (MockedStatic redisClientStatic = Mockito.mockStatic(RedisClient.class)) { + redisClientStatic.when(() -> RedisClient.create(any(RedisURI.class))).thenReturn(redisClient); + + notifier.start(); + notifier.start(); + + assertTrue(notifier.isRunning()); + verify(subConnection) + .addListener(org.mockito.ArgumentMatchers.>any()); + verify(pubSubCommands).subscribe("om:distributed-jobs:start", "om:distributed-jobs:complete"); + + notifier.stop(); + + assertFalse(notifier.isRunning()); + verify(pubSubCommands) + .unsubscribe("om:distributed-jobs:start", "om:distributed-jobs:complete"); + verify(subConnection).close(); + verify(pubConnection).close(); + verify(redisClient).shutdown(); + } + } + + @Test + void startRegistersListenerThatHandlesRemoteMessages() { + CacheConfig config = cacheConfig("redis://cache:6380"); + RedisJobNotifier notifier = new RedisJobNotifier(config, "server-123"); + RedisClient redisClient = mock(RedisClient.class); + StatefulRedisPubSubConnection subConnection = + mock(StatefulRedisPubSubConnection.class); + StatefulRedisConnection pubConnection = mock(StatefulRedisConnection.class); + RedisPubSubCommands pubSubCommands = mock(RedisPubSubCommands.class); + when(redisClient.connectPubSub()).thenReturn(subConnection); + when(redisClient.connect()).thenReturn(pubConnection); + when(subConnection.sync()).thenReturn(pubSubCommands); + + try (MockedStatic redisClientStatic = Mockito.mockStatic(RedisClient.class)) { + redisClientStatic.when(() -> RedisClient.create(any(RedisURI.class))).thenReturn(redisClient); + + notifier.start(); + AtomicReference callbackJob = new AtomicReference<>(); + notifier.onJobStarted(callbackJob::set); + + @SuppressWarnings("unchecked") + ArgumentCaptor> listenerCaptor = + ArgumentCaptor.forClass(RedisPubSubAdapter.class); + verify(subConnection).addListener(listenerCaptor.capture()); + + UUID jobId = UUID.randomUUID(); + listenerCaptor.getValue().message("om:distributed-jobs:start", jobId + "|SEARCH_INDEX|other"); + + assertEquals(jobId, callbackJob.get()); + } + } + + @Test + void startFailureResetsRunningState() { + RedisJobNotifier notifier = + new RedisJobNotifier(cacheConfig("redis://cache:6379"), "server-123"); + + try (MockedStatic redisClientStatic = Mockito.mockStatic(RedisClient.class)) { + redisClientStatic + .when(() -> RedisClient.create(any(RedisURI.class))) + .thenThrow(new IllegalStateException("redis down")); + + assertThrows(RuntimeException.class, notifier::start); + assertFalse(notifier.isRunning()); + } + } + + @Test + void stopReturnsWhenNotifierWasNeverStarted() { + RedisJobNotifier notifier = + new RedisJobNotifier(cacheConfig("redis://cache:6379"), "server-123"); + + notifier.stop(); + + assertFalse(notifier.isRunning()); + } + + @Test + void stopSwallowsShutdownExceptions() throws Exception { + RedisJobNotifier notifier = + new RedisJobNotifier(cacheConfig("redis://cache:6379"), "server-123"); + StatefulRedisPubSubConnection subConnection = + mock(StatefulRedisPubSubConnection.class); + RedisPubSubCommands pubSubCommands = mock(RedisPubSubCommands.class); + when(subConnection.sync()).thenReturn(pubSubCommands); + Mockito.doThrow(new IllegalStateException("unsubscribe failed")) + .when(pubSubCommands) + .unsubscribe("om:distributed-jobs:start", "om:distributed-jobs:complete"); + getRunningFlag(notifier).set(true); + setField(notifier, "subConnection", subConnection); + setField(notifier, "pubConnection", mock(StatefulRedisConnection.class)); + setField(notifier, "redisClient", mock(RedisClient.class)); + + notifier.stop(); + + assertFalse(notifier.isRunning()); + } + + @Test + void notifyMethodsAndInboundMessagesRespectSourceServer() throws Exception { + RedisJobNotifier notifier = new RedisJobNotifier(cacheConfig("cache:6379"), "server-123"); + StatefulRedisConnection pubConnection = mock(StatefulRedisConnection.class); + RedisCommands redisCommands = mock(RedisCommands.class); + when(pubConnection.sync()).thenReturn(redisCommands); + when(redisCommands.publish(any(), any())).thenReturn(2L); + getRunningFlag(notifier).set(true); + setField(notifier, "pubConnection", pubConnection); + + UUID jobId = UUID.randomUUID(); + notifier.notifyJobStarted(jobId, "SEARCH_INDEX"); + notifier.notifyJobCompleted(jobId); + + verify(redisCommands).publish("om:distributed-jobs:start", jobId + "|SEARCH_INDEX|server-123"); + verify(redisCommands).publish("om:distributed-jobs:complete", jobId + "|COMPLETED|server-123"); + + AtomicReference callbackJob = new AtomicReference<>(); + notifier.onJobStarted(callbackJob::set); + invokeHandleMessage( + notifier, "om:distributed-jobs:start", jobId + "|SEARCH_INDEX|other-server"); + assertEquals(jobId, callbackJob.get()); + + callbackJob.set(null); + invokeHandleMessage(notifier, "om:distributed-jobs:start", jobId + "|SEARCH_INDEX|server-123"); + assertNull(callbackJob.get()); + + invokeHandleMessage(notifier, "om:distributed-jobs:start", "invalid"); + invokeHandleMessage(notifier, "om:distributed-jobs:start", "not-a-uuid|SEARCH_INDEX|other"); + invokeHandleMessage( + notifier, "om:distributed-jobs:complete", jobId + "|COMPLETED|other-server"); + } + + @Test + void notifyMethodsSwallowPublishFailures() throws Exception { + RedisJobNotifier notifier = new RedisJobNotifier(cacheConfig("cache:6379"), "server-123"); + StatefulRedisConnection pubConnection = mock(StatefulRedisConnection.class); + RedisCommands redisCommands = mock(RedisCommands.class); + when(pubConnection.sync()).thenReturn(redisCommands); + when(redisCommands.publish(eq("om:distributed-jobs:start"), any())) + .thenThrow(new IllegalStateException("publish failed")); + when(redisCommands.publish(eq("om:distributed-jobs:complete"), any())) + .thenThrow(new IllegalStateException("publish failed")); + getRunningFlag(notifier).set(true); + setField(notifier, "pubConnection", pubConnection); + + notifier.notifyJobStarted(UUID.randomUUID(), "SEARCH_INDEX"); + notifier.notifyJobCompleted(UUID.randomUUID()); + } + + @Test + void notifyMethodsSkipWhenNotRunningOrWithoutPublisher() { + RedisJobNotifier notifier = new RedisJobNotifier(cacheConfig("cache:6379"), "server-123"); + StatefulRedisConnection pubConnection = mock(StatefulRedisConnection.class); + + notifier.notifyJobStarted(UUID.randomUUID(), "SEARCH_INDEX"); + notifier.notifyJobCompleted(UUID.randomUUID()); + verify(pubConnection, never()).sync(); + } + + @Test + void buildRedisUriSupportsUrlVariantsAndAuthentication() throws Exception { + CacheConfig config = cacheConfig("redis://cache.example.com:6380"); + config.redis.authType = CacheConfig.AuthType.PASSWORD; + config.redis.username = "user"; + config.redis.passwordRef = "secret"; + config.redis.useSSL = true; + config.redis.database = 4; + config.redis.connectTimeoutMs = 1234; + + RedisJobNotifier notifier = new RedisJobNotifier(config, "server-123"); + RedisURI uri = (RedisURI) invokePrivate(notifier, "buildRedisURI"); + + assertEquals("cache.example.com", uri.getHost()); + assertEquals(6380, uri.getPort()); + assertTrue(uri.isSsl()); + assertEquals(4, uri.getDatabase()); + assertEquals(Duration.ofMillis(1234), uri.getTimeout()); + assertEquals("user", uri.getUsername()); + + CacheConfig hostOnlyConfig = cacheConfig("redis-host"); + RedisURI hostOnlyUri = + (RedisURI) + invokePrivate(new RedisJobNotifier(hostOnlyConfig, "server-123"), "buildRedisURI"); + assertEquals("redis-host", hostOnlyUri.getHost()); + assertEquals(6379, hostOnlyUri.getPort()); + + CacheConfig hostPortConfig = cacheConfig("cache.example.com:6381"); + RedisURI hostPortUri = + (RedisURI) + invokePrivate(new RedisJobNotifier(hostPortConfig, "server-123"), "buildRedisURI"); + assertEquals("cache.example.com", hostPortUri.getHost()); + assertEquals(6381, hostPortUri.getPort()); + + CacheConfig passwordOnlyConfig = cacheConfig("cache.example.com:6382"); + passwordOnlyConfig.redis.authType = CacheConfig.AuthType.PASSWORD; + passwordOnlyConfig.redis.passwordRef = "secret"; + RedisURI passwordOnlyUri = + (RedisURI) + invokePrivate(new RedisJobNotifier(passwordOnlyConfig, "server-123"), "buildRedisURI"); + assertEquals("cache.example.com", passwordOnlyUri.getHost()); + assertEquals(6382, passwordOnlyUri.getPort()); + } + + @Test + void exposedTypeMatchesRedisImplementation() { + RedisJobNotifier notifier = new RedisJobNotifier(cacheConfig("cache:6379"), "server-123"); + + assertEquals("redis-pubsub", notifier.getType()); + } + + private void invokeHandleMessage(RedisJobNotifier notifier, String channel, String message) + throws Exception { + Method method = + notifier.getClass().getDeclaredMethod("handleMessage", String.class, String.class); + method.setAccessible(true); + method.invoke(notifier, channel, message); + } + + private Object invokePrivate(RedisJobNotifier notifier, String methodName) throws Exception { + Method method = notifier.getClass().getDeclaredMethod(methodName); + method.setAccessible(true); + return method.invoke(notifier); + } + + private AtomicBoolean getRunningFlag(RedisJobNotifier notifier) throws Exception { + return (AtomicBoolean) getField(notifier, "running"); + } + + private Object getField(Object target, String name) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + return field.get(target); + } + + private void setField(Object target, String name, Object value) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } + + private CacheConfig cacheConfig(String url) { + CacheConfig cacheConfig = new CacheConfig(); + cacheConfig.redis.url = url; + cacheConfig.redis.authType = CacheConfig.AuthType.NONE; + cacheConfig.redis.connectTimeoutMs = 2_000; + return cacheConfig; + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListenerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListenerTest.java index 1622b348cade..fb3bd4fb9bc5 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListenerTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/LoggingProgressListenerTest.java @@ -36,6 +36,8 @@ void onJobConfiguredInitializesLoggerAndTracksSettings() throws Exception { .maxConcurrentRequests(8) .payloadSize(2L * 1024 * 1024) .autoTune(true) + .recreateIndex(true) + .useDistributedIndexing(true) .build(); listener.onJobConfigured(context, config); @@ -51,7 +53,8 @@ void onJobConfiguredInitializesLoggerAndTracksSettings() throws Exception { assertEquals("8", details.get("Max Concurrent Requests")); assertEquals("2.0 MB", details.get("Payload Size")); assertEquals("Enabled", details.get("Auto-tune")); - assertEquals("Staged indexes with alias promotion", details.get("Indexing Mode")); + assertEquals("Yes", details.get("Recreate Index")); + assertEquals("Yes", details.get("Distributed Mode")); } @Test diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListenerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListenerTest.java index c4a4e7f9ef05..5653b9364723 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListenerTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/QuartzProgressListenerTest.java @@ -274,6 +274,8 @@ private ReindexingConfiguration configuration() { .queueSize(100) .maxConcurrentRequests(5) .payloadSize(4_096) + .recreateIndex(true) + .useDistributedIndexing(true) .build(); } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListenerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListenerTest.java index 2050ec63eb5f..63f8512cd10c 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListenerTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/bundles/searchIndex/listeners/SlackProgressListenerTest.java @@ -37,6 +37,7 @@ void jobConfiguredFormatsAllEntitiesAndPublishesConfiguration() throws Exception .maxConcurrentRequests(8) .payloadSize(5L * 1024 * 1024) .autoTune(true) + .recreateIndex(false) .build(); listener.onJobConfigured(mock(ReindexingJobContext.class), config); @@ -52,7 +53,7 @@ void jobConfiguredFormatsAllEntitiesAndPublishesConfiguration() throws Exception assertEquals("2", details.get("Producer threads")); assertEquals("500", details.get("Queue size")); assertEquals("1", details.get("Total entities")); - assertEquals("Staged indexes with alias promotion", details.get("Indexing mode")); + assertEquals("No", details.get("Recreating indices")); assertEquals("5 MB", details.get("Payload size")); assertEquals("8", details.get("Concurrent requests")); } @@ -73,6 +74,7 @@ void delegatesProgressCompletionAndErrorNotifications() throws Exception { .queueSize(200) .maxConcurrentRequests(3) .payloadSize(2L * 1024 * 1024) + .recreateIndex(true) .build(); Stats stats = new Stats() diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/apps/logging/AppRunLogAppenderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/apps/logging/AppRunLogAppenderTest.java index 5dc7d37e5e16..5a720aafdacf 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/apps/logging/AppRunLogAppenderTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/apps/logging/AppRunLogAppenderTest.java @@ -192,7 +192,7 @@ void concurrentWritesFromMultipleThreadsAreSafe() throws InterruptedException { @Test void formatLineProducesJsonMatchingDropwizardLayout() { LoggingEvent event = createEvent("reindex started", Map.of()); - event.setLoggerName("org.openmetadata.service.apps.bundles.searchIndex.ReindexingOrchestrator"); + event.setLoggerName("org.openmetadata.service.apps.bundles.searchIndex.SearchIndexExecutor"); event.setTimeStamp(1774260643332L); String line = AppRunLogAppender.formatLine(event); assertTrue(line.startsWith("{\"timestamp\":1774260643332,"), "should start with timestamp"); @@ -200,7 +200,7 @@ void formatLineProducesJsonMatchingDropwizardLayout() { assertTrue(line.contains("\"thread\":\"test-thread\""), "should contain thread"); assertTrue( line.contains( - "\"logger\":\"org.openmetadata.service.apps.bundles.searchIndex.ReindexingOrchestrator\""), + "\"logger\":\"org.openmetadata.service.apps.bundles.searchIndex.SearchIndexExecutor\""), "should contain full logger name"); assertTrue(line.contains("\"message\":\"reindex started\""), "should contain message"); assertTrue(line.endsWith("}"), "should be valid JSON object"); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/cache/EntityCacheBypassTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/cache/EntityCacheBypassTest.java index 0918211b63c4..76c43b1d07dc 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/cache/EntityCacheBypassTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/cache/EntityCacheBypassTest.java @@ -24,7 +24,8 @@ /** * Behaviour tests for the reindex cache-bypass thread-local. Pins the contract used by - * {@code PartitionWorker.processPartition} to opt reader threads out of the entity cache. + * {@code PartitionWorker.processPartition}, {@code EntityReader.readKeysetBatches}, and + * {@code SearchIndexExecutor} to opt their reader threads out of the entity cache. */ class EntityCacheBypassTest { diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfJsonLdContextTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfJsonLdContextTest.java index 0a4cd1dce0d5..2fb9452e8dd7 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfJsonLdContextTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfJsonLdContextTest.java @@ -19,6 +19,9 @@ class RdfJsonLdContextTest { private static ObjectMapper objectMapper; private static JsonNode baseContext; private static JsonNode lineageContext; + private static JsonNode governanceContext; + private static JsonNode aiContext; + private static JsonNode automationContext; @BeforeAll static void loadContexts() throws Exception { @@ -44,6 +47,42 @@ static void loadContexts() throws Exception { } } } + + try (InputStream is = + RdfJsonLdContextTest.class.getResourceAsStream("/rdf/contexts/governance.jsonld")) { + if (is != null) { + JsonNode contextDoc = objectMapper.readTree(is); + if (contextDoc.get("@context").isArray()) { + governanceContext = contextDoc.get("@context").get(1); + } else { + governanceContext = contextDoc.get("@context"); + } + } + } + + try (InputStream is = + RdfJsonLdContextTest.class.getResourceAsStream("/rdf/contexts/ai.jsonld")) { + if (is != null) { + JsonNode contextDoc = objectMapper.readTree(is); + if (contextDoc.get("@context").isArray()) { + aiContext = contextDoc.get("@context").get(1); + } else { + aiContext = contextDoc.get("@context"); + } + } + } + + try (InputStream is = + RdfJsonLdContextTest.class.getResourceAsStream("/rdf/contexts/automation.jsonld")) { + if (is != null) { + JsonNode contextDoc = objectMapper.readTree(is); + if (contextDoc.get("@context").isArray()) { + automationContext = contextDoc.get("@context").get(1); + } else { + automationContext = contextDoc.get("@context"); + } + } + } } @Nested @@ -386,4 +425,143 @@ void testColumnsLineageMapping() { } } } + + @Nested + @DisplayName("P1.8: Governance Context SKOS Hierarchy") + class GovernanceContextSkosTests { + + @Test + @DisplayName("glossary field should map to skos:inScheme on a glossary term") + void testGlossaryMapsToInScheme() { + assertNotNull(governanceContext, "governance.jsonld should be loaded"); + JsonNode glossary = governanceContext.get("glossary"); + assertNotNull(glossary, "'glossary' field mapping must be defined"); + assertEquals( + "skos:inScheme", + glossary.get("@id").asText(), + "GlossaryTerm.glossary should use SKOS inScheme, not the legacy om:belongsToGlossary"); + assertEquals("@id", glossary.get("@type").asText()); + } + + @Test + @DisplayName("classification field should map to skos:inScheme on a tag") + void testClassificationMapsToInScheme() { + assertNotNull(governanceContext); + JsonNode classification = governanceContext.get("classification"); + assertNotNull(classification, "'classification' field mapping must be defined"); + assertEquals( + "skos:inScheme", + classification.get("@id").asText(), + "Tag.classification should use SKOS inScheme to align with skos:ConceptScheme membership"); + } + + @Test + @DisplayName("parent field should map to skos:broader") + void testParentMapsToSkosBroader() { + assertNotNull(governanceContext); + JsonNode parent = governanceContext.get("parent"); + assertNotNull(parent, "'parent' field mapping must be defined"); + assertEquals("skos:broader", parent.get("@id").asText()); + assertEquals("@id", parent.get("@type").asText()); + } + + @Test + @DisplayName("children field should map to skos:narrower with @set container") + void testChildrenMapsToSkosNarrower() { + assertNotNull(governanceContext); + JsonNode children = governanceContext.get("children"); + assertNotNull( + children, + "'children' field mapping must be defined (it replaces the prior 'childTerms' alias which referenced a non-existent field)"); + assertEquals("skos:narrower", children.get("@id").asText()); + assertEquals("@id", children.get("@type").asText()); + assertEquals("@set", children.get("@container").asText()); + } + + @Test + @DisplayName("Stale childTerms alias should no longer be present") + void testNoLegacyChildTermsAlias() { + assertNotNull(governanceContext); + assertNull( + governanceContext.get("childTerms"), + "Legacy 'childTerms' alias must not coexist with 'children' — GlossaryTerm has no 'childTerms' field, so the alias would never fire"); + } + + @Test + @DisplayName("DataContract and Persona fields should be wired to ontology predicates") + void testDataContractAndPersonaFieldsWired() { + assertNotNull(governanceContext); + assertNotNull(governanceContext.get("contractStatus")); + assertEquals( + "om:contractStatus", governanceContext.get("contractStatus").get("@id").asText()); + assertNotNull(governanceContext.get("appliesTo")); + assertEquals("om:appliesToEntity", governanceContext.get("appliesTo").get("@id").asText()); + assertEquals("@id", governanceContext.get("appliesTo").get("@type").asText()); + assertNotNull(governanceContext.get("users")); + assertEquals("om:appliesToUser", governanceContext.get("users").get("@id").asText()); + } + } + + @Nested + @DisplayName("P1.5: AI / Automation Contexts") + class AiAutomationContextTests { + + @Test + @DisplayName("ai.jsonld should be loadable and define LLMModel + AIApplication types") + void testAiContextLoaded() { + assertNotNull(aiContext, "ai.jsonld should be on the classpath"); + assertEquals("om:LLMModel", aiContext.get("LLMModel").get("@id").asText()); + assertEquals("om:AIApplication", aiContext.get("AIApplication").get("@id").asText()); + assertEquals("om:McpServer", aiContext.get("McpServer").get("@id").asText()); + assertEquals("om:PromptTemplate", aiContext.get("PromptTemplate").get("@id").asText()); + } + + @Test + @DisplayName("AgentExecution and McpExecution should be PROV activities") + void testExecutionsAreProvActivities() { + assertNotNull(aiContext); + JsonNode agentExec = aiContext.get("AgentExecution"); + assertNotNull(agentExec); + assertTrue( + agentExec.get("@type").toString().contains("prov:Activity"), + "AgentExecution must be typed as prov:Activity for cross-system PROV traversal"); + JsonNode mcpExec = aiContext.get("McpExecution"); + assertTrue(mcpExec.get("@type").toString().contains("prov:Activity")); + } + + @Test + @DisplayName("LLMModel.trainingDatasets should map to om:hasTrainingDataset for AI lineage") + void testTrainingDatasetsMapping() { + assertNotNull(aiContext); + JsonNode trainingDatasets = aiContext.get("trainingDatasets"); + assertNotNull( + trainingDatasets, + "trainingDatasets must be wired so AI lineage queries can traverse model -> dataset"); + assertEquals("om:hasTrainingDataset", trainingDatasets.get("@id").asText()); + assertEquals("@id", trainingDatasets.get("@type").asText()); + assertEquals("@set", trainingDatasets.get("@container").asText()); + } + + @Test + @DisplayName("AIApplication.models should map to om:usesModel object property") + void testAiAppUsesModelMapping() { + assertNotNull(aiContext); + JsonNode models = aiContext.get("models"); + assertNotNull(models); + assertEquals("om:usesModel", models.get("@id").asText()); + assertEquals("@id", models.get("@type").asText()); + } + + @Test + @DisplayName("automation.jsonld should define Workflow + WorkflowInstance types") + void testAutomationContextLoaded() { + assertNotNull(automationContext, "automation.jsonld should be on the classpath"); + assertEquals("om:Workflow", automationContext.get("Workflow").get("@id").asText()); + JsonNode wfInstance = automationContext.get("WorkflowInstance"); + assertNotNull(wfInstance); + assertTrue( + wfInstance.get("@type").toString().contains("prov:Activity"), + "WorkflowInstance is a single run and must be typed as prov:Activity"); + } + } } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java index ca55bee53116..644fdc78a5e7 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/RdfPropertyMapperTest.java @@ -313,7 +313,7 @@ void testDownstreamEdges() throws Exception { } @Test - @DisplayName("Column lineage should be stored with fromColumn and toColumn properties") + @DisplayName("Column lineage should emit URI references plus FQN strings for back-compat") void testColumnLineage() throws Exception { ArrayNode upstreamEdges = objectMapper.createArrayNode(); ObjectNode edge = objectMapper.createObjectNode(); @@ -328,10 +328,10 @@ void testColumnLineage() throws Exception { ObjectNode colLineage = objectMapper.createObjectNode(); ArrayNode fromColumns = objectMapper.createArrayNode(); - fromColumns.add("source_table.column_a"); - fromColumns.add("source_table.column_b"); + fromColumns.add("service.db.schema.source_table.column_a"); + fromColumns.add("service.db.schema.source_table.column_b"); colLineage.set("fromColumns", fromColumns); - colLineage.put("toColumn", "target_table.merged_column"); + colLineage.put("toColumn", "service.db.schema.target_table.merged_column"); colLineage.put("function", "CONCAT(column_a, column_b)"); columnsLineage.add(colLineage); @@ -346,22 +346,437 @@ void testColumnLineage() throws Exception { method.setAccessible(true); method.invoke(propertyMapper, "upstreamEdges", upstreamEdges, entityResource, model); - // Find column lineage in the model Property hasColumnLineage = model.createProperty(OM_NS, "hasColumnLineage"); StmtIterator stmts = model.listStatements(null, hasColumnLineage, (Resource) null); assertTrue(stmts.hasNext(), "Should have column lineage"); + Resource colLineageResource = stmts.next().getObject().asResource(); + + Resource expectedFromA = + model.createResource( + RdfUtils.columnUri(BASE_URI, "service.db.schema.source_table.column_a")); + Resource expectedFromB = + model.createResource( + RdfUtils.columnUri(BASE_URI, "service.db.schema.source_table.column_b")); + Resource expectedTo = + model.createResource( + RdfUtils.columnUri(BASE_URI, "service.db.schema.target_table.merged_column")); + + Property fromColumn = model.createProperty(OM_NS, "fromColumn"); + Property toColumn = model.createProperty(OM_NS, "toColumn"); + assertTrue( + model.contains(colLineageResource, fromColumn, expectedFromA), + "fromColumn should reference URI for column_a"); + assertTrue( + model.contains(colLineageResource, fromColumn, expectedFromB), + "fromColumn should reference URI for column_b"); + assertTrue( + model.contains(colLineageResource, toColumn, expectedTo), + "toColumn should reference URI for merged_column"); + + Property fromColumnFqn = model.createProperty(OM_NS, "fromColumnFqn"); + Property toColumnFqn = model.createProperty(OM_NS, "toColumnFqn"); + assertTrue( + model.contains( + colLineageResource, fromColumnFqn, "service.db.schema.source_table.column_a"), + "fromColumnFqn literal should be retained for back-compat"); + assertTrue( + model.contains( + colLineageResource, toColumnFqn, "service.db.schema.target_table.merged_column"), + "toColumnFqn literal should be retained for back-compat"); + + Resource columnClass = model.createResource(OM_NS + "Column"); + assertTrue( + model.contains(expectedFromA, RDF.type, columnClass), + "Source column resource should be typed as om:Column"); + assertTrue( + model.contains(expectedTo, RDF.type, columnClass), + "Target column resource should be typed as om:Column"); + + Property transformFunc = model.createProperty(OM_NS, "transformFunction"); + assertTrue( + model.contains(colLineageResource, transformFunc, "CONCAT(column_a, column_b)"), + "transformFunction should be stored as a literal on the column-lineage resource"); + } + } + + @Nested + @DisplayName("P1.1: Column resource emission") + class ColumnResourceTests { + + @Test + @DisplayName("Table.columns should be emitted as named om:Column resources at FQN-derived URIs") + void testTableColumnsEmittedAsNamedResources() throws Exception { + Map contextCache = new HashMap<>(); + contextCache.put("dataAsset-complete", Map.of()); + propertyMapper = new RdfPropertyMapper(BASE_URI, objectMapper, contextCache); + + ArrayNode columns = objectMapper.createArrayNode(); + ObjectNode pkColumn = objectMapper.createObjectNode(); + pkColumn.put("name", "id"); + pkColumn.put("dataType", "BIGINT"); + pkColumn.put("constraint", "PRIMARY_KEY"); + pkColumn.put("ordinalPosition", 0); + pkColumn.put("description", "Primary key"); + pkColumn.put("fullyQualifiedName", "service.db.schema.orders.id"); + columns.add(pkColumn); + + ObjectNode amountColumn = objectMapper.createObjectNode(); + amountColumn.put("name", "amount"); + amountColumn.put("dataType", "DECIMAL"); + amountColumn.put("ordinalPosition", 1); + amountColumn.put("fullyQualifiedName", "service.db.schema.orders.amount"); + columns.add(amountColumn); + + invokePrivate( + "emitColumns", + new Class[] {JsonNode.class, Resource.class, Model.class}, + columns, + entityResource, + model); + + Resource pkResource = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.schema.orders.id")); + Resource amountResource = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.schema.orders.amount")); + + Property hasColumn = model.createProperty(OM_NS, "hasColumn"); + assertTrue( + model.contains(entityResource, hasColumn, pkResource), + "Table should link to PK column via om:hasColumn"); + assertTrue( + model.contains(entityResource, hasColumn, amountResource), + "Table should link to amount column via om:hasColumn"); + + Resource columnClass = model.createResource(OM_NS + "Column"); + assertTrue(model.contains(pkResource, RDF.type, columnClass)); + assertTrue( + model.contains(pkResource, model.createProperty(OM_NS, "columnDataType"), "BIGINT")); + assertTrue( + model.getProperty(pkResource, model.createProperty(OM_NS, "isPrimaryKey")).getBoolean(), + "Primary key constraint should set om:isPrimaryKey true"); + assertFalse( + model.getProperty(pkResource, model.createProperty(OM_NS, "isNullable")).getBoolean(), + "Primary key implies om:isNullable false"); + assertTrue( + model.contains(amountResource, model.createProperty(OM_NS, "columnDataType"), "DECIMAL")); + } + + @Test + @DisplayName("Nested struct/map columns should link via om:hasChildColumn") + void testNestedChildColumns() throws Exception { + ArrayNode columns = objectMapper.createArrayNode(); + ObjectNode struct = objectMapper.createObjectNode(); + struct.put("name", "address"); + struct.put("dataType", "STRUCT"); + struct.put("fullyQualifiedName", "service.db.schema.users.address"); + + ArrayNode children = objectMapper.createArrayNode(); + ObjectNode street = objectMapper.createObjectNode(); + street.put("name", "street"); + street.put("dataType", "VARCHAR"); + street.put("fullyQualifiedName", "service.db.schema.users.address.street"); + children.add(street); + struct.set("children", children); + + columns.add(struct); + + invokePrivate( + "emitColumns", + new Class[] {JsonNode.class, Resource.class, Model.class}, + columns, + entityResource, + model); + + Resource addressResource = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.schema.users.address")); + Resource streetResource = + model.createResource( + RdfUtils.columnUri(BASE_URI, "service.db.schema.users.address.street")); + + assertTrue( + model.contains( + addressResource, model.createProperty(OM_NS, "hasChildColumn"), streetResource), + "Parent struct column should link to child via om:hasChildColumn"); + assertTrue(model.contains(streetResource, RDF.type, model.createResource(OM_NS + "Column"))); + } + + @Test + @DisplayName("Per-column constraints map to isPrimaryKey, isNullable, and isUnique") + void testPerColumnConstraintFlags() throws Exception { + ArrayNode columns = objectMapper.createArrayNode(); + columns.add(columnNode("id", "BIGINT", "service.db.s.t.id", "PRIMARY_KEY")); + columns.add(columnNode("email", "VARCHAR", "service.db.s.t.email", "UNIQUE")); + columns.add(columnNode("country", "VARCHAR", "service.db.s.t.country", "NOT_NULL")); + columns.add(columnNode("nickname", "VARCHAR", "service.db.s.t.nickname", "NULL")); + + invokePrivate( + "emitColumns", + new Class[] {JsonNode.class, Resource.class, Model.class}, + columns, + entityResource, + model); + + Resource id = model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.t.id")); + Resource email = model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.t.email")); + Resource country = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.t.country")); + Resource nickname = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.t.nickname")); + + Property isPrimaryKey = model.createProperty(OM_NS, "isPrimaryKey"); + Property isUnique = model.createProperty(OM_NS, "isUnique"); + Property isNullable = model.createProperty(OM_NS, "isNullable"); + + assertTrue(model.getProperty(id, isPrimaryKey).getBoolean()); + assertTrue(model.getProperty(id, isUnique).getBoolean()); + assertFalse(model.getProperty(id, isNullable).getBoolean()); + + assertTrue(model.getProperty(email, isUnique).getBoolean()); + assertFalse( + model.contains(email, isPrimaryKey), + "UNIQUE alone should not imply primary-key membership"); + + assertFalse(model.getProperty(country, isNullable).getBoolean()); + assertTrue(model.getProperty(nickname, isNullable).getBoolean()); + } + + @Test + @DisplayName("FOREIGN_KEY table constraint emits om:references and TableConstraint resource") + void testForeignKeyTableConstraint() throws Exception { + ArrayNode constraints = objectMapper.createArrayNode(); + ObjectNode fk = objectMapper.createObjectNode(); + fk.put("constraintType", "FOREIGN_KEY"); + fk.put("relationshipType", "MANY_TO_ONE"); + ArrayNode cols = objectMapper.createArrayNode(); + cols.add("customer_id"); + fk.set("columns", cols); + ArrayNode referred = objectMapper.createArrayNode(); + referred.add("service.db.s.customers.id"); + fk.set("referredColumns", referred); + constraints.add(fk); + + invokePrivate( + "emitTableConstraints", + new Class[] {JsonNode.class, String.class, Resource.class, Model.class}, + constraints, + "service.db.s.orders", + entityResource, + model); - // Verify fromColumn is stored - Property fromColumnProp = model.createProperty(OM_NS, "fromColumn"); - assertTrue(model.contains(null, fromColumnProp), "Should have fromColumn properties"); + Resource customerIdCol = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.orders.customer_id")); + Resource referredCol = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.customers.id")); - // Verify toColumn is stored - Property toColumnProp = model.createProperty(OM_NS, "toColumn"); - assertTrue(model.contains(null, toColumnProp), "Should have toColumn property"); + Property references = model.createProperty(OM_NS, "references"); + assertTrue( + model.contains(customerIdCol, references, referredCol), + "FK should produce direct om:references triple between source and referred column"); + + Property hasConstraint = model.createProperty(OM_NS, "hasConstraint"); + Resource constraintResource = + model.listObjectsOfProperty(entityResource, hasConstraint).next().asResource(); + assertTrue( + model.contains( + constraintResource, RDF.type, model.createResource(OM_NS + "TableConstraint"))); + assertTrue( + model.contains( + constraintResource, model.createProperty(OM_NS, "constraintType"), "FOREIGN_KEY")); + assertTrue( + model.contains( + constraintResource, model.createProperty(OM_NS, "relationshipType"), "MANY_TO_ONE")); + assertTrue( + model.contains( + constraintResource, + model.createProperty(OM_NS, "hasConstrainedColumn"), + customerIdCol)); + assertTrue( + model.contains( + constraintResource, model.createProperty(OM_NS, "hasReferredColumn"), referredCol)); + } + + @Test + @DisplayName("Multi-column PRIMARY_KEY constraint marks every member column") + void testMultiColumnPrimaryKey() throws Exception { + ArrayNode constraints = objectMapper.createArrayNode(); + ObjectNode pk = objectMapper.createObjectNode(); + pk.put("constraintType", "PRIMARY_KEY"); + ArrayNode cols = objectMapper.createArrayNode(); + cols.add("tenant_id"); + cols.add("user_id"); + pk.set("columns", cols); + constraints.add(pk); + + invokePrivate( + "emitTableConstraints", + new Class[] {JsonNode.class, String.class, Resource.class, Model.class}, + constraints, + "service.db.s.users", + entityResource, + model); + + Resource tenantId = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.users.tenant_id")); + Resource userId = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.users.user_id")); + + Property isPrimaryKey = model.createProperty(OM_NS, "isPrimaryKey"); + assertTrue(model.getProperty(tenantId, isPrimaryKey).getBoolean()); + assertTrue(model.getProperty(userId, isPrimaryKey).getBoolean()); + } + + private ObjectNode columnNode(String name, String dataType, String fqn, String constraint) { + ObjectNode col = objectMapper.createObjectNode(); + col.put("name", name); + col.put("dataType", dataType); + col.put("fullyQualifiedName", fqn); + if (constraint != null) { + col.put("constraint", constraint); + } + return col; + } + + @Test + @DisplayName("Column.profile is emitted as DQV measurements rather than a JSON literal") + void testColumnProfileEmittedAsDqv() throws Exception { + ArrayNode columns = objectMapper.createArrayNode(); + ObjectNode col = objectMapper.createObjectNode(); + col.put("name", "email"); + col.put("dataType", "VARCHAR"); + col.put("fullyQualifiedName", "service.db.s.users.email"); + ObjectNode profile = objectMapper.createObjectNode(); + profile.put("valuesCount", 1000); + profile.put("nullCount", 12); + profile.put("nullProportion", 0.012); + profile.put("uniqueCount", 985); + profile.put("timestamp", 1714300000000L); + col.set("profile", profile); + columns.add(col); + + invokePrivate( + "emitColumns", + new Class[] {JsonNode.class, Resource.class, Model.class}, + columns, + entityResource, + model); + + Resource emailColumn = + model.createResource(RdfUtils.columnUri(BASE_URI, "service.db.s.users.email")); + Property hasMeasurement = + model.createProperty("http://www.w3.org/ns/dqv#", "hasQualityMeasurement"); + java.util.List measurements = + model.listObjectsOfProperty(emailColumn, hasMeasurement).toList().stream() + .map(node -> node.asResource()) + .toList(); + assertEquals( + 4, + measurements.size(), + "Expected 4 numeric profile metrics (valuesCount, nullCount, nullProportion, uniqueCount)"); + + Property isMeasurementOf = + model.createProperty("http://www.w3.org/ns/dqv#", "isMeasurementOf"); + Property dqvValue = model.createProperty("http://www.w3.org/ns/dqv#", "value"); + java.util.Map byMetric = new java.util.HashMap<>(); + for (Resource m : measurements) { + Resource metric = model.getProperty(m, isMeasurementOf).getObject().asResource(); + double v = model.getProperty(m, dqvValue).getDouble(); + byMetric.put(metric.getURI(), v); + } + assertEquals(1000.0, byMetric.get(OM_NS + "ValuesCountMetric"), 0.0); + assertEquals(12.0, byMetric.get(OM_NS + "NullCountMetric"), 0.0); + assertEquals(0.012, byMetric.get(OM_NS + "NullProportionMetric"), 1e-9); + assertEquals(985.0, byMetric.get(OM_NS + "UniqueCountMetric"), 0.0); + + // Each measurement should also be tied back to the column via dqv:computedOn. + Property computedOn = model.createProperty("http://www.w3.org/ns/dqv#", "computedOn"); + for (Resource m : measurements) { + assertTrue(model.contains(m, computedOn, emailColumn)); + } + } + + @Test + @DisplayName("Pipeline run is emitted as a prov:Activity tied to inputs and outputs") + void testPipelineRunEmitsProvActivity() throws Exception { + ObjectNode pipelineStatus = objectMapper.createObjectNode(); + pipelineStatus.put("timestamp", 1714300000000L); + pipelineStatus.put("endTime", 1714300120000L); + pipelineStatus.put("executionStatus", "Successful"); + pipelineStatus.put("executionId", "airflow-run-123"); + ArrayNode inputs = objectMapper.createArrayNode(); + ObjectNode in = objectMapper.createObjectNode(); + in.put("datasetFQN", "service.db.s.source"); + inputs.add(in); + pipelineStatus.set("inputs", inputs); + ArrayNode outputs = objectMapper.createArrayNode(); + ObjectNode out = objectMapper.createObjectNode(); + out.put("datasetFQN", "service.db.s.target"); + outputs.add(out); + pipelineStatus.set("outputs", outputs); + ObjectNode executedBy = objectMapper.createObjectNode(); + executedBy.put("id", UUID.randomUUID().toString()); + executedBy.put("type", "user"); + pipelineStatus.set("executedBy", executedBy); + + java.lang.reflect.Method method = + org.openmetadata.service.rdf.translator.RdfActivityMapper.class.getDeclaredMethod( + "emitPipelineActivity", + JsonNode.class, + String.class, + Resource.class, + String.class, + Model.class); + method.setAccessible(true); + method.invoke( + null, pipelineStatus, "service.pipeline.daily_etl", entityResource, BASE_URI, model); + + Property hasExecution = model.createProperty(OM_NS, "hasExecution"); + Resource activity = + model.listObjectsOfProperty(entityResource, hasExecution).next().asResource(); + + assertTrue( + model.contains( + activity, RDF.type, model.createResource("http://www.w3.org/ns/prov#Activity"))); + assertTrue( + model.contains(activity, model.createProperty(OM_NS, "executionStatus"), "Successful")); + assertTrue( + model.contains(activity, model.createProperty(OM_NS, "executionId"), "airflow-run-123")); + // PROV-O: activity-to-activity relation. Pipeline run wasInformedBy pipeline definition. + assertTrue( + model.contains( + activity, + model.createProperty("http://www.w3.org/ns/prov#", "wasInformedBy"), + entityResource)); + assertTrue( + model.contains( + activity, model.createProperty("http://www.w3.org/ns/prov#", "startedAtTime"))); + assertTrue( + model.contains( + activity, model.createProperty("http://www.w3.org/ns/prov#", "endedAtTime"))); + assertTrue( + model.contains(activity, model.createProperty("http://www.w3.org/ns/prov#", "used")), + "Activity should reference its input dataset via prov:used"); + assertTrue( + model.contains(activity, model.createProperty("http://www.w3.org/ns/prov#", "generated")), + "Activity should reference its output dataset via prov:generated"); + assertTrue( + model.contains( + activity, model.createProperty("http://www.w3.org/ns/prov#", "wasAssociatedWith")), + "Activity should record who triggered the run via prov:wasAssociatedWith"); + } + + @Test + @DisplayName("RdfUtils.columnUri should be deterministic and percent-encode FQNs") + void testColumnUri() { + String uri = RdfUtils.columnUri(BASE_URI, "service.db.schema.orders.amount"); + assertEquals(BASE_URI + "entity/column/service.db.schema.orders.amount", uri); + + String specialUri = RdfUtils.columnUri(BASE_URI, "service db.weird name"); + assertTrue( + specialUri.contains("service+db.weird+name") || specialUri.contains("service%20db"), + "FQN with whitespace should be percent-encoded"); - // Verify transformation function is stored - Property transformFuncProp = model.createProperty(OM_NS, "transformFunction"); - assertTrue(model.contains(null, transformFuncProp), "Should have transformFunction property"); + assertNull(RdfUtils.columnUri(BASE_URI, null)); + assertNull(RdfUtils.columnUri(BASE_URI, "")); } } diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/extension/CustomOntologyValidatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/extension/CustomOntologyValidatorTest.java new file mode 100644 index 000000000000..a6131b63a710 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/extension/CustomOntologyValidatorTest.java @@ -0,0 +1,345 @@ +package org.openmetadata.service.rdf.extension; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.openmetadata.schema.api.configuration.rdf.CustomOntology; +import org.openmetadata.schema.api.configuration.rdf.CustomOntologyClass; +import org.openmetadata.schema.api.configuration.rdf.CustomOntologyProperty; + +class CustomOntologyValidatorTest { + + private static final String EXT_NS = "https://open-metadata.org/ontology-extension/"; + + private static CustomOntologyClass cls(String localName, String... parents) { + return new CustomOntologyClass() + .withUri(EXT_NS + localName) + .withSubClassOf(java.util.List.of(parents)); + } + + private static CustomOntologyProperty objProp(String localName, String domain, String range) { + return new CustomOntologyProperty() + .withUri(EXT_NS + localName) + .withType(CustomOntologyProperty.Type.OBJECT_PROPERTY) + .withDomain(domain) + .withRange(range); + } + + private static CustomOntologyProperty datatypeProp( + String localName, String domain, String range) { + return new CustomOntologyProperty() + .withUri(EXT_NS + localName) + .withType(CustomOntologyProperty.Type.DATATYPE_PROPERTY) + .withDomain(domain) + .withRange(range); + } + + private static CustomOntology ext(String name) { + return new CustomOntology().withName(name); + } + + @Nested + @DisplayName("Required fields and shape") + class RequiredFieldsAndShape { + + @Test + @DisplayName("Null extension is rejected") + void nullExtension() { + assertTrue( + CustomOntologyValidator.validate(null).stream() + .anyMatch(e -> e.contains("must not be null"))); + } + + @Test + @DisplayName("Blank name is rejected") + void blankName() { + assertTrue( + CustomOntologyValidator.validate(ext("")).stream() + .anyMatch(e -> e.contains("'name' must not be blank"))); + } + + @Test + @DisplayName("Name with uppercase is rejected") + void uppercaseName() { + assertTrue( + CustomOntologyValidator.validate(ext("MyExtension")).stream() + .anyMatch(e -> e.contains("name") && e.contains("lowercase"))); + } + + @Test + @DisplayName("Extension with no classes and no properties is rejected") + void emptyExtension() { + List errors = CustomOntologyValidator.validate(ext("empty-ext")); + assertTrue(errors.stream().anyMatch(e -> e.contains("at least one class or property"))); + } + } + + @Nested + @DisplayName("Namespace enforcement") + class NamespaceEnforcement { + + @Test + @DisplayName("Class URI in canonical om: namespace is rejected (cannot redefine)") + void canonicalNamespaceClassRejected() { + CustomOntology e = + ext("redefine-table") + .withClasses( + List.of( + new CustomOntologyClass() + .withUri("https://open-metadata.org/ontology/Table") + .withSubClassOf(List.of("https://open-metadata.org/ontology/Entity")))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch( + err -> err.contains("om-extension namespace") && err.contains("ontology/Table"))); + } + + @Test + @DisplayName("Property URI outside om-extension namespace is rejected") + void canonicalNamespacePropertyRejected() { + CustomOntology e = + ext("bad-prop") + .withClasses(List.of(cls("Foo", "om:Entity"))) + .withProperties( + List.of( + new CustomOntologyProperty() + .withUri("https://example.org/somewhere") + .withType(CustomOntologyProperty.Type.OBJECT_PROPERTY) + .withDomain(EXT_NS + "Foo") + .withRange("om:Entity"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("om-extension namespace"))); + } + } + + @Nested + @DisplayName("Class hierarchy checks") + class ClassHierarchy { + + @Test + @DisplayName("Class without subClassOf is rejected") + void classNeedsParent() { + CustomOntologyClass orphan = new CustomOntologyClass().withUri(EXT_NS + "Orphan"); + CustomOntology e = ext("orphan-ext").withClasses(List.of(orphan)); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("must declare at least one subClassOf parent"))); + } + + @Test + @DisplayName("Class referencing unknown canonical parent is rejected") + void unknownCanonicalParentRejected() { + CustomOntology e = + ext("unknown-parent") + .withClasses( + List.of(cls("Widget", "https://open-metadata.org/ontology/NonexistentClass"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("unknown parent class"))); + } + + @Test + @DisplayName("Class referencing canonical om: short-form parent is accepted") + void shortFormCanonicalParent() { + CustomOntology e = ext("short-form").withClasses(List.of(cls("Widget", "om:DataAsset"))); + assertTrue(CustomOntologyValidator.validate(e).isEmpty()); + } + + @Test + @DisplayName("Class referencing another class in the same extension is accepted") + void siblingExtensionClassReference() { + CustomOntology e = + ext("siblings") + .withClasses(List.of(cls("Parent", "om:Entity"), cls("Child", EXT_NS + "Parent"))); + List errors = CustomOntologyValidator.validate(e); + assertTrue(errors.isEmpty(), "Got: " + errors); + } + + @Test + @DisplayName("Cycle in class hierarchy is detected (A → B → A)") + void hierarchyCycleDetected() { + CustomOntology e = + ext("cycle").withClasses(List.of(cls("A", EXT_NS + "B"), cls("B", EXT_NS + "A"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("contains a cycle")), + "Validator must detect 2-node hierarchy cycles"); + } + + @Test + @DisplayName("Self-referencing class is detected as a cycle") + void selfCycleDetected() { + CustomOntology e = ext("self-cycle").withClasses(List.of(cls("Self", EXT_NS + "Self"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("contains a cycle"))); + } + + @Test + @DisplayName("3-node cycle (A → B → C → A) is detected") + void threeNodeCycleDetected() { + CustomOntology e = + ext("3-cycle") + .withClasses( + List.of(cls("A", EXT_NS + "B"), cls("B", EXT_NS + "C"), cls("C", EXT_NS + "A"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("contains a cycle"))); + } + + @Test + @DisplayName("Duplicate class URI within the same extension is rejected") + void duplicateClassUriRejected() { + CustomOntology e = + ext("dupes").withClasses(List.of(cls("Same", "om:Entity"), cls("Same", "om:Entity"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("duplicate class URI"))); + } + } + + @Nested + @DisplayName("Property checks") + class PropertyChecks { + + @Test + @DisplayName("ObjectProperty with unknown domain is rejected") + void unknownDomainRejected() { + CustomOntology e = + ext("unknown-domain") + .withClasses(List.of(cls("Foo", "om:Entity"))) + .withProperties( + List.of( + objProp( + "rel", "https://open-metadata.org/ontology/NoSuchClass", "om:Entity"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("domain") && err.contains("not a known"))); + } + + @Test + @DisplayName("ObjectProperty with non-class range is rejected") + void objectPropertyRangeMustBeClass() { + CustomOntology e = + ext("bad-range") + .withClasses(List.of(cls("Foo", "om:Entity"))) + .withProperties( + List.of( + objProp("rel", EXT_NS + "Foo", "http://www.w3.org/2001/XMLSchema#string"))); + // ObjectProperty range that is an xsd type is rejected because it's not a known class. + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("range") && err.contains("not a known"))); + } + + @Test + @DisplayName("DatatypeProperty range must be an xsd: datatype URI") + void datatypePropertyRequiresXsdRange() { + CustomOntology e = + ext("bad-dt-range") + .withClasses(List.of(cls("Foo", "om:Entity"))) + .withProperties(List.of(datatypeProp("score", EXT_NS + "Foo", EXT_NS + "Foo"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch( + err -> err.contains("DatatypeProperty") && err.contains("xsd: datatype URI"))); + } + + @Test + @DisplayName("Valid DatatypeProperty with xsd:string range is accepted") + void validDatatypeProperty() { + CustomOntology e = + ext("valid-dt") + .withClasses(List.of(cls("Foo", "om:Entity"))) + .withProperties( + List.of( + datatypeProp( + "score", EXT_NS + "Foo", "http://www.w3.org/2001/XMLSchema#string"))); + List errors = CustomOntologyValidator.validate(e); + assertTrue(errors.isEmpty(), "Got: " + errors); + } + + @Test + @DisplayName("Duplicate property URI within the same extension is rejected") + void duplicatePropertyUri() { + CustomOntology e = + ext("dupe-props") + .withClasses(List.of(cls("Foo", "om:Entity"))) + .withProperties( + List.of( + objProp("rel", EXT_NS + "Foo", "om:Entity"), + objProp("rel", EXT_NS + "Foo", "om:Entity"))); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("duplicate property URI"))); + } + + @Test + @DisplayName("Property with no domain is rejected") + void missingDomain() { + CustomOntologyProperty p = + new CustomOntologyProperty() + .withUri(EXT_NS + "rel") + .withType(CustomOntologyProperty.Type.OBJECT_PROPERTY) + .withRange("om:Entity"); + CustomOntology e = + ext("no-domain").withClasses(List.of(cls("Foo", "om:Entity"))).withProperties(List.of(p)); + assertTrue( + CustomOntologyValidator.validate(e).stream() + .anyMatch(err -> err.contains("missing 'domain'"))); + } + } + + @Nested + @DisplayName("Happy path") + class HappyPath { + + @Test + @DisplayName("Full valid extension passes validation cleanly") + void fullValidExtension() { + CustomOntology e = + ext("regulatory-controls") + .withDescription("SOX compliance controls") + .withClasses( + List.of( + cls("RegulatoryControl", "om:Entity"), + cls("SoxControl", EXT_NS + "RegulatoryControl"))) + .withProperties( + List.of( + objProp("hasControl", "om:DataAsset", EXT_NS + "RegulatoryControl"), + datatypeProp( + "controlOwnerEmail", + EXT_NS + "RegulatoryControl", + "http://www.w3.org/2001/XMLSchema#string"))); + List errors = CustomOntologyValidator.validate(e); + assertTrue(errors.isEmpty(), "Expected no errors but got: " + errors); + assertTrue(CustomOntologyValidator.isValid(e)); + } + + @Test + @DisplayName("Property-only extension (no classes) is accepted") + void propertyOnlyExtension() { + CustomOntology e = + ext("annotations") + .withProperties( + List.of( + datatypeProp( + "soxRelevant", + "om:DataAsset", + "http://www.w3.org/2001/XMLSchema#boolean"))); + assertTrue(CustomOntologyValidator.validate(e).isEmpty()); + } + } + + @Test + @DisplayName("isValid wrapper returns boolean and logs on failure") + void isValidWrapper() { + assertFalse(CustomOntologyValidator.isValid(null)); + assertFalse(CustomOntologyValidator.isValid(ext("empty-test"))); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/federation/SparqlFederationGuardTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/federation/SparqlFederationGuardTest.java new file mode 100644 index 000000000000..f6db29b8bc76 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/federation/SparqlFederationGuardTest.java @@ -0,0 +1,253 @@ +package org.openmetadata.service.rdf.federation; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import java.util.Set; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +/** + * Failure-mode coverage for the SPARQL federation guard. Each test names exactly the case it + * is exercising — these are the queries that real users (and adversaries) will send. + */ +class SparqlFederationGuardTest { + + private static final String WIKIDATA = "https://query.wikidata.org/sparql"; + private static final String DBPEDIA = "https://dbpedia.org/sparql"; + + private SparqlFederationGuard disabled() { + return new SparqlFederationGuard(false, Set.of()); + } + + private SparqlFederationGuard withAllowlist(String... endpoints) { + return new SparqlFederationGuard(true, Set.of(endpoints)); + } + + @Nested + @DisplayName("Federation disabled (default policy)") + class FederationDisabled { + + @Test + @DisplayName("Plain query without SERVICE is always allowed") + void plainQueryAllowed() { + String q = "SELECT * WHERE { ?s ?p ?o } LIMIT 1"; + assertTrue(disabled().firstDisallowedEndpoint(q).isEmpty()); + } + + @Test + @DisplayName("Any SERVICE clause is rejected when federation is disabled") + void serviceRejectedWhenDisabled() { + String q = "SELECT * WHERE { SERVICE <" + WIKIDATA + "> { ?s ?p ?o } } LIMIT 1"; + assertEquals(WIKIDATA, disabled().firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("enforce throws FederationDisallowedException with helpful message") + void enforceThrows() { + String q = "SELECT * WHERE { SERVICE <" + WIKIDATA + "> { ?s ?p ?o } }"; + var ex = + assertThrows( + SparqlFederationGuard.FederationDisallowedException.class, + () -> disabled().enforce(q)); + assertEquals(WIKIDATA, ex.getBlockedEndpoint()); + assertFalse(ex.isFederationEnabled()); + assertTrue(ex.getMessage().contains("federated SPARQL is disabled")); + } + } + + @Nested + @DisplayName("Federation enabled with allowlist") + class FederationEnabled { + + @Test + @DisplayName("Allowlisted endpoint passes") + void allowlistedPasses() { + String q = "SELECT * WHERE { SERVICE <" + WIKIDATA + "> { ?s ?p ?o } }"; + assertTrue(withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).isEmpty()); + } + + @Test + @DisplayName("Endpoint not on allowlist is rejected even when federation is enabled") + void notAllowlistedRejected() { + String q = "SELECT * WHERE { SERVICE <" + DBPEDIA + "> { ?s ?p ?o } }"; + assertEquals(DBPEDIA, withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("Multiple SERVICE clauses: allowed + disallowed → reject the disallowed one") + void mixedServicesRejected() { + String q = + "SELECT * WHERE { " + + " SERVICE <" + + WIKIDATA + + "> { ?s ?p ?o } " + + " SERVICE <" + + DBPEDIA + + "> { ?s ?p ?o } " + + "}"; + assertEquals(DBPEDIA, withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("All SERVICE clauses allowlisted → query allowed") + void allServicesAllowed() { + String q = + "SELECT * WHERE { " + + " SERVICE <" + + WIKIDATA + + "> { ?a ?b ?c } " + + " SERVICE <" + + DBPEDIA + + "> { ?d ?e ?f } " + + "}"; + assertTrue(withAllowlist(WIKIDATA, DBPEDIA).firstDisallowedEndpoint(q).isEmpty()); + } + + @Test + @DisplayName("SERVICE SILENT is detected the same as SERVICE") + void silentServiceDetected() { + String q = "SELECT * WHERE { SERVICE SILENT <" + DBPEDIA + "> { ?s ?p ?o } }"; + assertEquals(DBPEDIA, withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("SERVICE with variable endpoint is always rejected — can't be allowlisted") + void variableServiceRejected() { + String q = "SELECT * WHERE { ?endpoint a SERVICE ?endpoint { ?s ?p ?o } }"; + // The SPARQL parser may or may not accept this exact form depending on context; if it does, + // a variable endpoint is unprovable against a static allowlist and must be rejected. + var blocked = withAllowlist(WIKIDATA).firstDisallowedEndpoint(q); + // If the engine parses it, we expect rejection; if not, no SERVICE was extracted, which is + // also acceptable because the engine itself will reject the query. + blocked.ifPresent(b -> assertTrue(b.startsWith("?"))); + } + + @Test + @DisplayName("Trailing-slash mismatch: allowlist must match the URI exactly") + void trailingSlashDoesNotMatch() { + String q = "SELECT * WHERE { SERVICE <" + WIKIDATA + "/> { ?s ?p ?o } }"; + assertEquals( + WIKIDATA + "/", + withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).orElseThrow(), + "We compare endpoint URIs as strings, including trailing slashes; this is documented behavior"); + } + + @Test + @DisplayName("Nested SERVICE inside OPTIONAL is detected") + void nestedInsideOptional() { + String q = "SELECT * WHERE { ?s ?p ?o OPTIONAL { SERVICE <" + DBPEDIA + "> { ?s ?p ?o } } }"; + assertEquals(DBPEDIA, withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("Nested SERVICE inside UNION branch is detected") + void nestedInsideUnion() { + String q = "SELECT * WHERE { { ?s ?p ?o } UNION { SERVICE <" + DBPEDIA + "> { ?s ?p ?o } } }"; + assertEquals(DBPEDIA, withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("SERVICE inside subquery is detected") + void nestedInsideSubquery() { + String q = + "SELECT * WHERE { " + + " { SELECT ?s WHERE { SERVICE <" + + DBPEDIA + + "> { ?s ?p ?o } } } " + + "}"; + assertEquals(DBPEDIA, withAllowlist(WIKIDATA).firstDisallowedEndpoint(q).orElseThrow()); + } + } + + @Nested + @DisplayName("Adversarial inputs") + class AdversarialInputs { + + @Test + @DisplayName("The literal text 'SERVICE' inside a string literal must NOT trigger the guard") + void serviceLiteralInString() { + String q = "SELECT * WHERE { ?s ?p \"SERVICE <" + DBPEDIA + ">\" }"; + assertTrue( + disabled().firstDisallowedEndpoint(q).isEmpty(), + "Regex-based detectors fail this; the parser-based guard must not"); + } + + @Test + @DisplayName("SPARQL comment with 'SERVICE' must not trigger the guard") + void serviceInComment() { + String q = "# SERVICE <" + DBPEDIA + "> { ?s ?p ?o }\n" + "SELECT * WHERE { ?s ?p ?o }"; + assertTrue(disabled().firstDisallowedEndpoint(q).isEmpty()); + } + + @Test + @DisplayName("Unparseable garbage SPARQL is passed through (engine emits its own parse error)") + void unparseableQueryPassesThrough() { + String garbage = "this is not sparql {{{}}}"; + assertTrue( + disabled().firstDisallowedEndpoint(garbage).isEmpty(), + "Guard must not turn a parse error into a federation error — engine handles parsing"); + } + + @Test + @DisplayName("Empty / null / whitespace queries are passed through") + void emptyQueriesPassedThrough() { + assertTrue(disabled().firstDisallowedEndpoint("").isEmpty()); + assertTrue(disabled().firstDisallowedEndpoint(" ").isEmpty()); + } + + @Test + @DisplayName("Lowercase 'service' keyword is detected (SPARQL is case-insensitive)") + void lowercaseService() { + String q = "select * where { service <" + DBPEDIA + "> { ?s ?p ?o } }"; + assertEquals(DBPEDIA, disabled().firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("ASK with SERVICE is also guarded (not just SELECT)") + void askQueryGuarded() { + String q = "ASK { SERVICE <" + DBPEDIA + "> { ?s ?p ?o } }"; + assertEquals(DBPEDIA, disabled().firstDisallowedEndpoint(q).orElseThrow()); + } + + @Test + @DisplayName("CONSTRUCT with SERVICE is also guarded") + void constructQueryGuarded() { + String q = "CONSTRUCT { ?s ?p ?o } WHERE { SERVICE <" + DBPEDIA + "> { ?s ?p ?o } }"; + assertEquals(DBPEDIA, disabled().firstDisallowedEndpoint(q).orElseThrow()); + } + } + + @Nested + @DisplayName("serviceEndpoints listing") + class ServiceEndpointsListing { + + @Test + @DisplayName("Returns endpoints in order of first appearance, deduplicated") + void listingOrderAndDedup() { + String q = + "SELECT * WHERE { " + + " SERVICE <" + + DBPEDIA + + "> { ?s ?p ?o } " + + " SERVICE <" + + WIKIDATA + + "> { ?s ?p ?o } " + + " SERVICE <" + + DBPEDIA + + "> { ?s ?p ?o } " + + "}"; + assertEquals(List.of(DBPEDIA, WIKIDATA), disabled().serviceEndpoints(q)); + } + + @Test + @DisplayName("Returns empty list for queries without SERVICE") + void emptyListForNoService() { + assertTrue(disabled().serviceEndpoints("SELECT ?s WHERE { ?s a ?t }").isEmpty()); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/inference/InferenceRuleValidatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/inference/InferenceRuleValidatorTest.java new file mode 100644 index 000000000000..a141a87932f6 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/inference/InferenceRuleValidatorTest.java @@ -0,0 +1,226 @@ +package org.openmetadata.service.rdf.inference; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.openmetadata.schema.api.configuration.rdf.InferenceRule; + +/** + * Failure-mode coverage for {@link InferenceRuleValidator}. Each test exercises a way the + * validator must reject a hostile or malformed rule before it can be applied server-side. + */ +class InferenceRuleValidatorTest { + + private static final String VALID_BODY = + "PREFIX om: \n" + + "CONSTRUCT { ?x om:transitivelyDerivedFrom ?y }\n" + + "WHERE { ?x + ?y . FILTER(?x != ?y) }"; + + private static InferenceRule rule(String name, String body) { + return new InferenceRule() + .withName(name) + .withRuleType(InferenceRule.RuleType.CONSTRUCT) + .withRuleBody(body) + .withEnabled(true) + .withPriority(100); + } + + private static List validate(InferenceRule rule) { + return InferenceRuleValidator.validate(rule); + } + + @Nested + @DisplayName("Required fields") + class RequiredFields { + + @Test + @DisplayName("Null rule produces a single error") + void nullRuleRejected() { + List errors = validate(null); + assertFalse(errors.isEmpty()); + assertTrue(errors.get(0).contains("must not be null")); + } + + @Test + @DisplayName("Blank name is rejected") + void blankNameRejected() { + assertTrue( + validate(rule("", VALID_BODY)).stream() + .anyMatch(e -> e.contains("'name' must not be blank"))); + } + + @Test + @DisplayName("Name with uppercase or spaces is rejected (must match the schema pattern)") + void invalidNamePatternRejected() { + assertTrue( + validate(rule("MyRule", VALID_BODY)).stream() + .anyMatch(e -> e.contains("name") && e.contains("lowercase"))); + assertTrue( + validate(rule("my rule", VALID_BODY)).stream() + .anyMatch(e -> e.contains("name") && e.contains("lowercase"))); + assertTrue( + validate(rule("a", VALID_BODY)).stream() + .anyMatch(e -> e.contains("name") && e.contains("3-64"))); + } + + @Test + @DisplayName("Blank ruleBody is rejected") + void blankBodyRejected() { + assertTrue( + validate(rule("good-name", "")).stream() + .anyMatch(e -> e.contains("'ruleBody' must not be blank"))); + } + } + + @Nested + @DisplayName("RuleType handling") + class RuleTypeHandling { + + @Test + @DisplayName("RDFS ruleType is rejected (reserved for future)") + void rdfsTypeRejected() { + InferenceRule r = rule("future-rdfs", VALID_BODY).withRuleType(InferenceRule.RuleType.RDFS); + assertTrue(validate(r).stream().anyMatch(e -> e.contains("RDFS is reserved for future use"))); + } + } + + @Nested + @DisplayName("SPARQL body shape") + class SparqlBodyShape { + + @Test + @DisplayName("Garbage SPARQL is rejected with a parse error") + void parseError() { + List errors = validate(rule("bad-syntax", "this is not sparql")); + assertTrue(errors.stream().anyMatch(e -> e.startsWith("ruleBody failed to parse as SPARQL"))); + } + + @Test + @DisplayName("SELECT body is rejected — must be CONSTRUCT for ruleType=CONSTRUCT") + void selectBodyRejected() { + String selectBody = "SELECT ?s WHERE { ?s ?p ?o }"; + assertTrue( + validate(rule("must-be-construct", selectBody)).stream() + .anyMatch(e -> e.contains("must be a SPARQL CONSTRUCT query"))); + } + + @Test + @DisplayName("ASK body is rejected") + void askBodyRejected() { + String askBody = "ASK { ?s ?p ?o }"; + assertTrue( + validate(rule("ask-not-allowed", askBody)).stream() + .anyMatch(e -> e.contains("must be a SPARQL CONSTRUCT"))); + } + + @Test + @DisplayName("DESCRIBE body is rejected") + void describeBodyRejected() { + String describeBody = "DESCRIBE "; + assertTrue( + validate(rule("describe-not-allowed", describeBody)).stream() + .anyMatch(e -> e.contains("must be a SPARQL CONSTRUCT"))); + } + + @Test + @DisplayName("CONSTRUCT with empty WHERE pattern is rejected") + void emptyWherePatternRejected() { + String emptyWhere = "CONSTRUCT { } WHERE { }"; + List errors = validate(rule("empty-where", emptyWhere)); + assertTrue( + errors.stream().anyMatch(e -> e.contains("non-empty WHERE pattern")), "Got: " + errors); + } + + @Test + @DisplayName("CONSTRUCT with empty template is rejected") + void emptyTemplateRejected() { + String emptyTemplate = "CONSTRUCT { } WHERE { ?s ?p ?o }"; + assertTrue( + validate(rule("empty-template", emptyTemplate)).stream() + .anyMatch(e -> e.contains("CONSTRUCT template must contain at least one triple"))); + } + } + + @Nested + @DisplayName("SERVICE rejection") + class ServiceRejection { + + @Test + @DisplayName("CONSTRUCT body containing a SERVICE clause is rejected") + void serviceClauseRejected() { + String body = + "CONSTRUCT { ?s ?p ?o } WHERE { SERVICE { ?s ?p ?o } }"; + List errors = validate(rule("federated-rule", body)); + assertTrue( + errors.stream().anyMatch(e -> e.contains("must not contain SERVICE clauses")), + "Got: " + errors); + } + + @Test + @DisplayName("SERVICE inside a subquery is also rejected") + void serviceInSubqueryRejected() { + String body = + "CONSTRUCT { ?s ?p ?o } WHERE { { SELECT ?s ?p ?o WHERE { SERVICE { ?s ?p ?o } } } }"; + assertTrue( + validate(rule("nested-service", body)).stream() + .anyMatch(e -> e.contains("must not contain SERVICE clauses"))); + } + } + + @Nested + @DisplayName("Priority bounds") + class PriorityBounds { + + @Test + @DisplayName("Priority below 0 is rejected") + void negativePriorityRejected() { + InferenceRule r = rule("low-pri", VALID_BODY).withPriority(-1); + assertTrue( + validate(r).stream().anyMatch(e -> e.contains("'priority' must be between 0 and 10000"))); + } + + @Test + @DisplayName("Priority above 10000 is rejected") + void highPriorityRejected() { + InferenceRule r = rule("high-pri", VALID_BODY).withPriority(10_001); + assertTrue(validate(r).stream().anyMatch(e -> e.contains("'priority' must be between"))); + } + } + + @Nested + @DisplayName("Happy path") + class HappyPath { + + @Test + @DisplayName("Well-formed CONSTRUCT rule passes validation with no errors") + void validRulePasses() { + InferenceRule r = rule("transitive-lineage", VALID_BODY); + List errors = validate(r); + assertTrue(errors.isEmpty(), "Expected no errors but got: " + errors); + } + } + + @Nested + @DisplayName("Starter pack rules ship valid") + class StarterPackValidation { + + @Test + @DisplayName("All shipped starter-pack rules pass validation") + void starterPackValid() { + InferenceRuleRegistry registry = InferenceRuleRegistry.getInstance(); + registry.loadStarterPackIfNeeded(); + List rules = registry.list(); + assertTrue(rules.size() >= 4, "Starter pack must ship at least 4 rules"); + for (InferenceRule rule : rules) { + List errors = validate(rule); + assertTrue( + errors.isEmpty(), + "Starter pack rule '" + rule.getName() + "' must validate cleanly. Got: " + errors); + } + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CentralityComputationTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CentralityComputationTest.java new file mode 100644 index 000000000000..babf3a63a6e3 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CentralityComputationTest.java @@ -0,0 +1,231 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.contains; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.Map; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.openmetadata.service.rdf.RdfRepository; + +class CentralityComputationTest { + + @Nested + @DisplayName("Predicate weights") + class PredicateWeights { + + @Test + @DisplayName("Lineage weighted highest, hasColumn weakest, unknown predicates excluded") + void weights() { + assertEquals(1.0, CentralityComputation.weightFor("prov:wasDerivedFrom")); + assertEquals(0.5, CentralityComputation.weightFor("om:hasTag")); + assertEquals(0.5, CentralityComputation.weightFor("om:hasGlossaryTerm")); + assertEquals(0.2, CentralityComputation.weightFor("om:hasColumn")); + assertEquals(0.0, CentralityComputation.weightFor("om:somethingElse")); + } + } + + @Nested + @DisplayName("SPARQL JSON → adjacency map parsing") + class GraphParsing { + + private static String row(String from, String to, String pred) { + return "{\"from\":{\"value\":\"" + + from + + "\"},\"to\":{\"value\":\"" + + to + + "\"},\"predicate\":{\"value\":\"" + + pred + + "\"}}"; + } + + private static String body(String... rows) { + return "{\"results\":{\"bindings\":[" + String.join(",", rows) + "]}}"; + } + + @Test + @DisplayName("Empty / null / blank input returns an empty graph") + void emptyInputs() { + assertTrue(CentralityComputation.parseGraph(null).isEmpty()); + assertTrue(CentralityComputation.parseGraph("").isEmpty()); + assertTrue(CentralityComputation.parseGraph("not json").isEmpty()); + assertTrue(CentralityComputation.parseGraph("{\"results\":{\"bindings\":[]}}").isEmpty()); + } + + @Test + @DisplayName("Single lineage edge produces weight 1.0") + void singleLineageEdge() { + Map> g = + CentralityComputation.parseGraph(body(row("urn:a", "urn:b", "prov:wasDerivedFrom"))); + assertEquals(1.0, g.get("urn:a").get("urn:b")); + } + + @Test + @DisplayName("hasColumn edge produces weight 0.2") + void hasColumnEdge() { + Map> g = + CentralityComputation.parseGraph(body(row("urn:t", "urn:c", "om:hasColumn"))); + assertEquals(0.2, g.get("urn:t").get("urn:c")); + } + + @Test + @DisplayName("Multiple edges to the same target sum their weights") + void parallelEdges() { + Map> g = + CentralityComputation.parseGraph( + body( + row("urn:a", "urn:b", "prov:wasDerivedFrom"), + row("urn:a", "urn:b", "om:hasColumn"))); + assertEquals(1.2, g.get("urn:a").get("urn:b"), 1e-9); + } + + @Test + @DisplayName("Unknown predicate produces weight 0.0 (effectively dropped)") + void unknownPredicateSilentlyDropped() { + Map> g = + CentralityComputation.parseGraph(body(row("urn:a", "urn:b", "om:unknown"))); + assertEquals(0.0, g.get("urn:a").get("urn:b")); + } + + @Test + @DisplayName("Rows missing 'from' or 'to' are skipped") + void missingBindingsSkipped() { + String partial = "{\"results\":{\"bindings\":[{\"from\":{\"value\":\"urn:a\"}}]}}"; + assertTrue(CentralityComputation.parseGraph(partial).isEmpty()); + } + } + + @Nested + @DisplayName("End-to-end computeAndPersist") + class EndToEnd { + + @Test + @DisplayName("Empty graph: returns 0 nodes, no SPARQL UPDATE") + void emptyGraph() { + RdfRepository repo = mock(RdfRepository.class); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenReturn("{\"results\":{\"bindings\":[]}}"); + CentralityComputation comp = new CentralityComputation(repo); + + CentralityComputation.Result r = comp.computeAndPersist("table"); + + assertEquals("table", r.entityType()); + assertEquals(0, r.nodesScored()); + assertFalse(r.converged()); + verify(repo, never()).executeSparqlUpdate(anyString()); + } + + @Test + @DisplayName("Lineage triangle persisted with normalized scores summing to 1.0") + void lineageTriangle() { + RdfRepository repo = mock(RdfRepository.class); + String body = + "{\"results\":{\"bindings\":[" + + "{\"from\":{\"value\":\"urn:a\"},\"to\":{\"value\":\"urn:b\"},\"predicate\":{\"value\":\"prov:wasDerivedFrom\"}}," + + "{\"from\":{\"value\":\"urn:b\"},\"to\":{\"value\":\"urn:c\"},\"predicate\":{\"value\":\"prov:wasDerivedFrom\"}}," + + "{\"from\":{\"value\":\"urn:c\"},\"to\":{\"value\":\"urn:a\"},\"predicate\":{\"value\":\"prov:wasDerivedFrom\"}}" + + "]}}"; + when(repo.executeSparqlQuery(anyString(), anyString())).thenReturn(body); + + CentralityComputation comp = new CentralityComputation(repo); + CentralityComputation.Result r = comp.computeAndPersist("table"); + + assertEquals(3, r.nodesScored()); + assertTrue(r.converged()); + + ArgumentCaptor update = ArgumentCaptor.forClass(String.class); + verify(repo, times(1)).executeSparqlUpdate(update.capture()); + String sparqlUpdate = update.getValue(); + assertTrue( + sparqlUpdate.contains( + "WITH "), + "Should write to the entityType-specific named graph: " + sparqlUpdate); + assertTrue(sparqlUpdate.contains("om:centralityScore")); + assertTrue(sparqlUpdate.contains("om:centralityRank")); + assertTrue(sparqlUpdate.contains("DELETE")); + assertTrue(sparqlUpdate.contains("INSERT DATA")); + // All three nodes referenced + assertTrue(sparqlUpdate.contains("urn:a")); + assertTrue(sparqlUpdate.contains("urn:b")); + assertTrue(sparqlUpdate.contains("urn:c")); + } + + @Test + @DisplayName("Star topology: hub gets the highest persisted score (rank 1)") + void hubGetsRank1() { + RdfRepository repo = mock(RdfRepository.class); + StringBuilder body = new StringBuilder("{\"results\":{\"bindings\":["); + for (int i = 0; i < 5; i++) { + if (i > 0) body.append(","); + body.append("{\"from\":{\"value\":\"urn:leaf-") + .append(i) + .append( + "\"},\"to\":{\"value\":\"urn:hub\"},\"predicate\":{\"value\":\"prov:wasDerivedFrom\"}}"); + } + body.append("]}}"); + when(repo.executeSparqlQuery(anyString(), anyString())).thenReturn(body.toString()); + + CentralityComputation comp = new CentralityComputation(repo); + comp.computeAndPersist("table"); + + ArgumentCaptor update = ArgumentCaptor.forClass(String.class); + verify(repo).executeSparqlUpdate(update.capture()); + String sparql = update.getValue(); + // The hub should appear before any leaf in the INSERT block (sorted by score desc). + int hubIdx = sparql.indexOf("urn:hub"); + int firstLeafIdx = sparql.indexOf("urn:leaf-"); + assertTrue(hubIdx > 0 && firstLeafIdx > 0); + assertTrue(hubIdx < firstLeafIdx, "Hub should be persisted first (highest score)"); + } + + @Test + @DisplayName("Bad entityType is rejected before any SPARQL is sent") + void badEntityTypeRejected() { + RdfRepository repo = mock(RdfRepository.class); + CentralityComputation comp = new CentralityComputation(repo); + + org.junit.jupiter.api.Assertions.assertThrows( + IllegalArgumentException.class, () -> comp.computeAndPersist("table OR 1=1")); + verify(repo, never()).executeSparqlQuery(anyString(), anyString()); + verify(repo, never()).executeSparqlUpdate(anyString()); + } + + @Test + @DisplayName("Repository SPARQL throws → empty result, no update attempted") + void repositoryThrows() { + RdfRepository repo = mock(RdfRepository.class); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenThrow(new RuntimeException("Fuseki unreachable")); + CentralityComputation comp = new CentralityComputation(repo); + + CentralityComputation.Result r = comp.computeAndPersist("table"); + assertEquals(0, r.nodesScored()); + verify(repo, never()).executeSparqlUpdate(anyString()); + } + + @Test + @DisplayName("Persisted graph URI uses lowercase entityType") + void graphUriLowercase() { + RdfRepository repo = mock(RdfRepository.class); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenReturn( + "{\"results\":{\"bindings\":[" + + "{\"from\":{\"value\":\"urn:a\"},\"to\":{\"value\":\"urn:b\"},\"predicate\":{\"value\":\"prov:wasDerivedFrom\"}}" + + "]}}"); + new CentralityComputation(repo).computeAndPersist("Dashboard"); + verify(repo) + .executeSparqlUpdate( + contains("")); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CoOccurrenceQueryBuilderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CoOccurrenceQueryBuilderTest.java new file mode 100644 index 000000000000..a47c7809dd86 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CoOccurrenceQueryBuilderTest.java @@ -0,0 +1,130 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.apache.jena.query.Query; +import org.apache.jena.query.QueryFactory; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +class CoOccurrenceQueryBuilderTest { + + @Nested + @DisplayName("Tag co-occurrence") + class TagCoOccurrence { + + @Test + @DisplayName("Generated query parses with Jena and selects the expected vars") + void parses() { + Query q = QueryFactory.create(CoOccurrenceQueryBuilder.tagCoOccurrence(2, 20)); + assertTrue(q.isSelectType()); + assertTrue(q.getResultVars().contains("tagA")); + assertTrue(q.getResultVars().contains("tagB")); + assertTrue(q.getResultVars().contains("count")); + } + + @Test + @DisplayName("Pairs are canonicalised so each pair appears once") + void canonicalised() { + String sparql = CoOccurrenceQueryBuilder.tagCoOccurrence(2, 20); + assertTrue( + sparql.contains("STR(?tagA) < STR(?tagB)"), + "Without canonical pair filter (a, b) and (b, a) would each appear"); + } + + @Test + @DisplayName("HAVING enforces the minimum count") + void havingPresent() { + String sparql = CoOccurrenceQueryBuilder.tagCoOccurrence(5, 20); + assertTrue(sparql.contains("HAVING (COUNT(?entity) >= 5)")); + } + + @Test + @DisplayName("minCount below 1 is clamped to 1") + void minCountClampedLow() { + String sparql = CoOccurrenceQueryBuilder.tagCoOccurrence(-100, 10); + assertTrue(sparql.contains("HAVING (COUNT(?entity) >= 1)")); + } + + @Test + @DisplayName("limit above MAX_LIMIT is clamped down") + void limitClampedHigh() { + String sparql = CoOccurrenceQueryBuilder.tagCoOccurrence(2, 9999); + assertTrue(sparql.endsWith("LIMIT " + CoOccurrenceQueryBuilder.MAX_LIMIT)); + } + + @Test + @DisplayName("limit below 1 is clamped to 1") + void limitClampedLow() { + String sparql = CoOccurrenceQueryBuilder.tagCoOccurrence(2, 0); + assertTrue(sparql.endsWith("LIMIT 1")); + } + + @Test + @DisplayName("ORDER BY DESC(?count) is present") + void orderByCount() { + assertTrue(CoOccurrenceQueryBuilder.tagCoOccurrence(2, 20).contains("ORDER BY DESC(?count)")); + } + } + + @Nested + @DisplayName("Glossary reach") + class GlossaryReach { + + @Test + @DisplayName("Generated query parses with Jena and counts DISTINCT domains") + void parses() { + Query q = QueryFactory.create(CoOccurrenceQueryBuilder.glossaryReach(2, 20)); + assertTrue(q.isSelectType()); + assertTrue(q.getResultVars().contains("term")); + assertTrue(q.getResultVars().contains("domainCount")); + String sparql = CoOccurrenceQueryBuilder.glossaryReach(2, 20); + assertTrue(sparql.contains("COUNT(DISTINCT ?domain)")); + } + + @Test + @DisplayName("HAVING enforces minDomains floor; clamped values are used") + void minDomainsClamped() { + assertTrue( + CoOccurrenceQueryBuilder.glossaryReach(-5, 20) + .contains("HAVING (COUNT(DISTINCT ?domain) >= 1)")); + assertTrue( + CoOccurrenceQueryBuilder.glossaryReach(7, 20) + .contains("HAVING (COUNT(DISTINCT ?domain) >= 7)")); + } + + @Test + @DisplayName("Joins om:hasGlossaryTerm with om:hasDomain on the same entity") + void joinsCorrectly() { + String sparql = CoOccurrenceQueryBuilder.glossaryReach(2, 20); + assertTrue(sparql.contains("?entity om:hasGlossaryTerm ?term")); + assertTrue(sparql.contains("?entity om:hasDomain ?domain")); + } + } + + @Nested + @DisplayName("Tag popularity") + class TagPopularity { + + @Test + @DisplayName("Generated query parses, counts DISTINCT entities, no HAVING required") + void parses() { + Query q = QueryFactory.create(CoOccurrenceQueryBuilder.tagPopularity(20)); + assertTrue(q.isSelectType()); + assertTrue(q.getResultVars().contains("tag")); + assertTrue(q.getResultVars().contains("entityCount")); + String sparql = CoOccurrenceQueryBuilder.tagPopularity(20); + assertTrue(sparql.contains("COUNT(DISTINCT ?entity)")); + } + + @Test + @DisplayName("limit clamping behaves like the other builders") + void clamping() { + assertTrue(CoOccurrenceQueryBuilder.tagPopularity(0).endsWith("LIMIT 1")); + assertTrue( + CoOccurrenceQueryBuilder.tagPopularity(9999) + .endsWith("LIMIT " + CoOccurrenceQueryBuilder.MAX_LIMIT)); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CommunityComputationTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CommunityComputationTest.java new file mode 100644 index 000000000000..21dfcdbed92b --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/CommunityComputationTest.java @@ -0,0 +1,306 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.contains; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.Map; +import org.apache.jena.query.QueryFactory; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.openmetadata.service.rdf.RdfRepository; + +class CommunityComputationTest { + + @Nested + @DisplayName("GraphType parsing") + class GraphTypeParsing { + + @Test + @DisplayName("Defaults to lineage when null/blank") + void defaults() { + assertEquals( + CommunityComputation.GraphType.LINEAGE, CommunityComputation.GraphType.parse(null)); + assertEquals( + CommunityComputation.GraphType.LINEAGE, CommunityComputation.GraphType.parse("")); + assertEquals( + CommunityComputation.GraphType.LINEAGE, CommunityComputation.GraphType.parse(" ")); + } + + @Test + @DisplayName("Aliases for tagCoOccurrence are accepted") + void tagAliases() { + assertEquals( + CommunityComputation.GraphType.TAG_CO_OCCURRENCE, + CommunityComputation.GraphType.parse("tagCoOccurrence")); + assertEquals( + CommunityComputation.GraphType.TAG_CO_OCCURRENCE, + CommunityComputation.GraphType.parse("tags")); + assertEquals( + CommunityComputation.GraphType.TAG_CO_OCCURRENCE, + CommunityComputation.GraphType.parse("TAG")); + assertEquals( + CommunityComputation.GraphType.TAG_CO_OCCURRENCE, + CommunityComputation.GraphType.parse("tag-co-occurrence")); + } + + @Test + @DisplayName("Unknown values are rejected") + void unknown() { + assertThrows( + IllegalArgumentException.class, + () -> CommunityComputation.GraphType.parse("citation_network")); + } + } + + @Nested + @DisplayName("SPARQL well-formedness") + class WellFormed { + + @Test + @DisplayName("Lineage graph SPARQL parses with Jena and uses both upstream/downstream edges") + void lineageSparqlParses() { + String sparql = CommunityComputation.lineageGraphSparql("Table"); + QueryFactory.create(sparql); + assertTrue(sparql.contains("prov:wasDerivedFrom")); + assertTrue(sparql.contains("om:upstream")); + assertTrue(sparql.contains("om:downstream")); + assertTrue(sparql.contains("a om:Table")); + } + + @Test + @DisplayName("Tag co-occurrence SPARQL parses and groups by pair with COUNT(?shared)") + void tagSparqlParses() { + String sparql = CommunityComputation.tagCoOccurrenceSparql("Dashboard"); + QueryFactory.create(sparql); + assertTrue(sparql.contains("om:hasTag")); + assertTrue(sparql.contains("om:hasGlossaryTerm")); + assertTrue(sparql.contains("COUNT(?shared)")); + assertTrue( + sparql.contains("STR(?from) < STR(?to)"), + "Pairs must be canonicalized to avoid double-counting"); + } + + @Test + @DisplayName("Listing SPARQL targets the entityType-specific named graph") + void listingSparqlTargetsNamedGraph() { + String sparql = CommunityComputation.listingSparql("table", "lineage"); + QueryFactory.create(sparql); + assertTrue( + sparql.contains( + "FROM ")); + assertTrue(sparql.contains("ORDER BY DESC(?size)")); + } + + @Test + @DisplayName("Listing SPARQL rejects bad entity / graph type before emitting SPARQL") + void listingSparqlValidatesInput() { + assertThrows( + IllegalArgumentException.class, + () -> CommunityComputation.listingSparql("table OR 1=1", "lineage")); + assertThrows( + IllegalArgumentException.class, + () -> CommunityComputation.listingSparql("table", "elsewhere")); + } + } + + @Nested + @DisplayName("SPARQL JSON → adjacency parsing") + class GraphParsing { + + private static String row(String from, String to, String weight) { + String w = weight == null ? "" : ",\"weight\":{\"value\":\"" + weight + "\"}"; + return "{\"from\":{\"value\":\"" + from + "\"},\"to\":{\"value\":\"" + to + "\"}" + w + "}"; + } + + private static String body(String... rows) { + return "{\"results\":{\"bindings\":[" + String.join(",", rows) + "]}}"; + } + + @Test + @DisplayName("Empty / null / blank input → empty adjacency") + void empty() { + assertTrue(CommunityComputation.parseGraph(null).isEmpty()); + assertTrue(CommunityComputation.parseGraph("").isEmpty()); + assertTrue(CommunityComputation.parseGraph("not json").isEmpty()); + assertTrue(CommunityComputation.parseGraph("{\"results\":{\"bindings\":[]}}").isEmpty()); + } + + @Test + @DisplayName("Single edge yields one directed entry (Louvain symmetrizes internally)") + void directedSingleEdge() { + Map> g = + CommunityComputation.parseGraph(body(row("urn:a", "urn:b", "1.0"))); + assertEquals(1.0, g.get("urn:a").get("urn:b")); + // The target node is registered (so Louvain sees it) but the reverse weight is NOT added — + // Louvain.addAllEdges adds both directions to its internal adjacency. Duplicating here + // would double-count every edge weight. + assertTrue(g.containsKey("urn:b"), "target node must be a key so Louvain enumerates it"); + assertTrue( + g.get("urn:b") == null || !g.get("urn:b").containsKey("urn:a"), + "reverse-direction weight must not be populated by parseGraph"); + } + + @Test + @DisplayName("Missing weight defaults to 1.0") + void weightDefaults() { + Map> g = + CommunityComputation.parseGraph(body(row("urn:a", "urn:b", null))); + assertEquals(1.0, g.get("urn:a").get("urn:b")); + } + + @Test + @DisplayName("Self-loops are dropped") + void selfLoopsDropped() { + Map> g = + CommunityComputation.parseGraph(body(row("urn:a", "urn:a", "1.0"))); + assertTrue(g.isEmpty()); + } + + @Test + @DisplayName("Non-positive weights are dropped") + void nonPositiveDropped() { + Map> g = + CommunityComputation.parseGraph( + body(row("urn:a", "urn:b", "0"), row("urn:c", "urn:d", "-1"))); + assertTrue(g.isEmpty()); + } + + @Test + @DisplayName("Non-numeric weights fall back to default and the edge is kept") + void nonNumericFallback() { + Map> g = + CommunityComputation.parseGraph(body(row("urn:a", "urn:b", "garbage"))); + assertEquals(1.0, g.get("urn:a").get("urn:b")); + } + } + + @Nested + @DisplayName("End-to-end computeAndPersist") + class EndToEnd { + + @Test + @DisplayName("Empty graph: returns 0 communities, no SPARQL UPDATE") + void emptyGraph() { + RdfRepository repo = mock(RdfRepository.class); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenReturn("{\"results\":{\"bindings\":[]}}"); + CommunityComputation comp = new CommunityComputation(repo); + CommunityComputation.Result r = comp.computeAndPersist("table", "lineage"); + assertEquals(0, r.communities()); + assertEquals(0, r.membersTotal()); + verify(repo, never()).executeSparqlUpdate(anyString()); + } + + @Test + @DisplayName("Two cliques persist as two om:Community resources with correct members") + void twoCliques() { + RdfRepository repo = mock(RdfRepository.class); + String body = + "{\"results\":{\"bindings\":[" + + edge("urn:a", "urn:b") + + "," + + edge("urn:b", "urn:c") + + "," + + edge("urn:a", "urn:c") + + "," + + edge("urn:x", "urn:y") + + "," + + edge("urn:y", "urn:z") + + "," + + edge("urn:x", "urn:z") + + "," + + edge("urn:c", "urn:x") + + "]}}"; + when(repo.executeSparqlQuery(anyString(), anyString())).thenReturn(body); + + CommunityComputation.Result r = + new CommunityComputation(repo).computeAndPersist("table", "lineage"); + assertEquals(2, r.communities()); + assertEquals(6, r.membersTotal()); + + ArgumentCaptor update = ArgumentCaptor.forClass(String.class); + verify(repo, times(1)).executeSparqlUpdate(update.capture()); + String sparql = update.getValue(); + assertTrue(sparql.contains("om:Community")); + assertTrue(sparql.contains("om:hasMember")); + assertTrue(sparql.contains("om:communitySize")); + assertTrue(sparql.contains("om:modularity")); + assertTrue( + sparql.contains( + "WITH "), + "Must persist into lineage/table named graph: " + sparql); + assertTrue(sparql.contains("urn:a")); + assertTrue(sparql.contains("urn:z")); + } + + @Test + @DisplayName("Bad entity type is rejected before any SPARQL is sent") + void badEntityTypeRejected() { + RdfRepository repo = mock(RdfRepository.class); + assertThrows( + IllegalArgumentException.class, + () -> new CommunityComputation(repo).computeAndPersist("foo OR 1=1", "lineage")); + verify(repo, never()).executeSparqlQuery(anyString(), anyString()); + } + + @Test + @DisplayName("Unknown graphType is rejected before any SPARQL is sent") + void badGraphType() { + RdfRepository repo = mock(RdfRepository.class); + assertThrows( + IllegalArgumentException.class, + () -> new CommunityComputation(repo).computeAndPersist("table", "weather")); + verify(repo, never()).executeSparqlQuery(anyString(), anyString()); + } + + @Test + @DisplayName("SPARQL exception during extraction → empty result, no update") + void sparqlError() { + RdfRepository repo = mock(RdfRepository.class); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenThrow(new RuntimeException("Fuseki down")); + CommunityComputation.Result r = + new CommunityComputation(repo).computeAndPersist("table", "lineage"); + assertEquals(0, r.communities()); + verify(repo, never()).executeSparqlUpdate(anyString()); + } + + @Test + @DisplayName("Tag-co-occurrence run uses the tagCoOccurrence named graph") + void tagsTargetsTagGraph() { + RdfRepository repo = mock(RdfRepository.class); + String body = "{\"results\":{\"bindings\":[" + edgeWithWeight("urn:a", "urn:b", 3.0) + "]}}"; + when(repo.executeSparqlQuery(anyString(), anyString())).thenReturn(body); + + new CommunityComputation(repo).computeAndPersist("table", "tagCoOccurrence"); + verify(repo) + .executeSparqlUpdate( + contains( + "")); + } + + private static String edge(String from, String to) { + return "{\"from\":{\"value\":\"" + from + "\"},\"to\":{\"value\":\"" + to + "\"}}"; + } + + private static String edgeWithWeight(String from, String to, double w) { + return "{\"from\":{\"value\":\"" + + from + + "\"},\"to\":{\"value\":\"" + + to + + "\"},\"weight\":{\"value\":\"" + + w + + "\"}}"; + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/ImportanceQueryBuilderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/ImportanceQueryBuilderTest.java new file mode 100644 index 000000000000..b10aca7ed5ce --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/ImportanceQueryBuilderTest.java @@ -0,0 +1,209 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.apache.jena.query.Query; +import org.apache.jena.query.QueryFactory; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +/** + * Tests for {@link ImportanceQueryBuilder}. Covers: + * + *
    + *
  1. Input validation — bad entityType / window / limit must be rejected, not sanitized + * silently. + *
  2. SPARQL well-formedness — every produced query must parse with Jena (catches typos in + * string interpolation). + *
  3. Score formula — the projected SELECT contains the right predicates and the BIND for + * score blends usage and downstream count with the right weights. + *
+ */ +class ImportanceQueryBuilderTest { + + @Nested + @DisplayName("Input validation") + class InputValidation { + + @Test + @DisplayName("Null entityType is rejected") + void nullEntityType() { + assertThrows( + IllegalArgumentException.class, () -> ImportanceQueryBuilder.build(null, "daily", 10)); + } + + @Test + @DisplayName("Blank entityType is rejected") + void blankEntityType() { + assertThrows( + IllegalArgumentException.class, () -> ImportanceQueryBuilder.build(" ", "daily", 10)); + } + + @Test + @DisplayName("Non-alphanumeric entityType is rejected (defends against SPARQL injection)") + void injectionAttemptRejected() { + assertThrows( + IllegalArgumentException.class, + () -> ImportanceQueryBuilder.build("table> ; DROP --", "daily", 10)); + assertThrows( + IllegalArgumentException.class, + () -> ImportanceQueryBuilder.build("table OR 1=1", "daily", 10)); + assertThrows( + IllegalArgumentException.class, + () -> ImportanceQueryBuilder.build("ta'ble", "daily", 10)); + } + + @Test + @DisplayName("Unknown window is rejected") + void unknownWindow() { + assertThrows( + IllegalArgumentException.class, + () -> ImportanceQueryBuilder.build("table", "yearly", 10)); + } + + @Test + @DisplayName("Window is normalized to lowercase") + void windowLowercased() { + String q = ImportanceQueryBuilder.build("table", "DAILY", 10); + assertTrue(q.contains("usageDailyPercentile")); + } + + @Test + @DisplayName("Empty / null window defaults to daily") + void windowDefaults() { + String q1 = ImportanceQueryBuilder.build("table", null, 10); + String q2 = ImportanceQueryBuilder.build("table", "", 10); + assertTrue(q1.contains("usageDailyPercentile")); + assertTrue(q2.contains("usageDailyPercentile")); + } + + @Test + @DisplayName("Limit below 1 is clamped to 1") + void limitClampedLow() { + String q = ImportanceQueryBuilder.build("table", "daily", -5); + assertTrue(q.endsWith("LIMIT 1"), "Got: " + q); + } + + @Test + @DisplayName("Limit above 100 is clamped to 100") + void limitClampedHigh() { + String q = ImportanceQueryBuilder.build("table", "daily", 9999); + assertTrue(q.endsWith("LIMIT 100")); + } + + @Test + @DisplayName("Default limit is 20") + void defaultLimit() { + assertEquals(20, ImportanceQueryBuilder.DEFAULT_LIMIT); + } + } + + @Nested + @DisplayName("Class-name capitalization") + class ClassCapitalization { + + @Test + @DisplayName("entityType is capitalized to match the OWL class") + void capitalize() { + assertTrue(ImportanceQueryBuilder.build("table", "daily", 10).contains("a om:Table")); + assertTrue(ImportanceQueryBuilder.build("dashboard", "daily", 10).contains("a om:Dashboard")); + assertTrue(ImportanceQueryBuilder.build("pipeline", "daily", 10).contains("a om:Pipeline")); + } + + @Test + @DisplayName("Already-capitalized entityType stays capitalized") + void alreadyCapitalized() { + assertTrue(ImportanceQueryBuilder.build("Table", "daily", 10).contains("a om:Table")); + } + } + + @Nested + @DisplayName("SPARQL well-formedness") + class WellFormed { + + @Test + @DisplayName("Generated query parses as a valid SPARQL Query (Jena)") + void parsesWithJena() { + Query q = QueryFactory.create(ImportanceQueryBuilder.build("table", "daily", 20)); + assertTrue(q.isSelectType(), "Expected SELECT query"); + assertTrue(q.getResultVars().contains("entity")); + assertTrue(q.getResultVars().contains("score")); + assertTrue(q.getResultVars().contains("usagePct")); + assertTrue(q.getResultVars().contains("downstreamCount")); + } + + @Test + @DisplayName("Generated query uses ORDER BY DESC(score)") + void orderByScore() { + String body = ImportanceQueryBuilder.build("table", "daily", 20); + assertTrue(body.contains("ORDER BY DESC(?score)"), "Got: " + body); + } + } + + @Nested + @DisplayName("Score formula") + class ScoreFormula { + + @Test + @DisplayName("Score blends usage (0.6) and downstream (0.4) with centrality reserved at 0.0") + void weightsExpressed() { + String body = ImportanceQueryBuilder.build("table", "daily", 20); + assertTrue(body.contains("0.6 * ?usageNorm"), "Got: " + body); + assertTrue(body.contains("0.4 * ?downstreamNorm"), "Got: " + body); + assertTrue( + body.contains("0.0 * ?centralityNorm"), + "Centrality term must be reserved at 0.0 until 3.1.b ships its PageRank fallback"); + } + + @Test + @DisplayName("Usage percentile is divided by 100 to land in 0-1") + void usageNormalizedTo01() { + assertTrue( + ImportanceQueryBuilder.build("table", "daily", 20) + .contains("COALESCE(?usagePct, 0.0) / 100.0")); + } + + @Test + @DisplayName("Downstream count is normalized by max-downstream-count via subquery") + void downstreamNormalizedByMax() { + String body = ImportanceQueryBuilder.build("table", "daily", 20); + assertTrue(body.contains("MAX(?dc) AS ?maxDownstream")); + assertTrue(body.contains("xsd:double(?downstreamCount) / xsd:double(?maxDownstream)")); + assertTrue( + body.contains("IF(?maxDownstream > 0,"), + "Must guard against division by zero when no entity has downstream lineage"); + } + } + + @Nested + @DisplayName("Window selection") + class WindowSelection { + + @Test + @DisplayName("daily window queries usageDailyPercentile") + void daily() { + assertTrue( + ImportanceQueryBuilder.build("table", "daily", 20) + .contains("om:usageDailyPercentile ?usagePct")); + } + + @Test + @DisplayName("weekly window queries usageWeeklyPercentile") + void weekly() { + assertTrue( + ImportanceQueryBuilder.build("table", "weekly", 20) + .contains("om:usageWeeklyPercentile ?usagePct")); + } + + @Test + @DisplayName("monthly window queries usageMonthlyPercentile") + void monthly() { + assertTrue( + ImportanceQueryBuilder.build("table", "monthly", 20) + .contains("om:usageMonthlyPercentile ?usagePct")); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LineagePathBuilderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LineagePathBuilderTest.java new file mode 100644 index 000000000000..c082f13fbcaf --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LineagePathBuilderTest.java @@ -0,0 +1,274 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import org.apache.jena.query.Query; +import org.apache.jena.query.QueryFactory; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +/** + * Tests for {@link LineagePathBuilder}. Covers: + * + *
    + *
  1. Input validation — bad URI / hop / direction values must throw, not silently sanitize. + *
  2. SPARQL well-formedness — every produced query must parse with Jena. + *
  3. Predicate UNIONs change with direction so the caller can swap walk modes without + * branching elsewhere. + *
+ */ +class LineagePathBuilderTest { + + @Nested + @DisplayName("URI validation") + class UriValidation { + + @Test + @DisplayName("Null URI is rejected") + void nullUri() { + assertThrows( + IllegalArgumentException.class, () -> LineagePathBuilder.validateNodeUri("from", null)); + } + + @Test + @DisplayName("Blank URI is rejected") + void blankUri() { + assertThrows( + IllegalArgumentException.class, () -> LineagePathBuilder.validateNodeUri("from", " ")); + } + + @Test + @DisplayName("Relative URI is rejected") + void relativeUri() { + assertThrows( + IllegalArgumentException.class, + () -> LineagePathBuilder.validateNodeUri("from", "/relative/path")); + } + + @Test + @DisplayName("Non-http scheme is rejected") + void nonHttpScheme() { + assertThrows( + IllegalArgumentException.class, + () -> LineagePathBuilder.validateNodeUri("from", "file:///etc/passwd")); + assertThrows( + IllegalArgumentException.class, + () -> LineagePathBuilder.validateNodeUri("from", "ftp://example.com/x")); + } + + @Test + @DisplayName("URI with angle bracket is rejected (defense against SPARQL injection)") + void angleBracketRejected() { + assertThrows( + IllegalArgumentException.class, + () -> + LineagePathBuilder.validateNodeUri( + "from", "https://example.com/x> ; DROP GRAPH LineagePathBuilder.validateNodeUri("from", "https://x.com/y\nINSERT")); + } + + @Test + @DisplayName("Malformed URI is rejected") + void malformedUri() { + assertThrows( + IllegalArgumentException.class, + () -> LineagePathBuilder.validateNodeUri("from", "https://[badhost")); + } + + @Test + @DisplayName("Valid http(s) URIs are accepted and trimmed") + void valid() { + assertEquals( + "https://open-metadata.org/instance/Table/abc", + LineagePathBuilder.validateNodeUri( + "from", " https://open-metadata.org/instance/Table/abc ")); + assertEquals( + "http://localhost:3030/foo", + LineagePathBuilder.validateNodeUri("from", "http://localhost:3030/foo")); + } + } + + @Nested + @DisplayName("Direction parsing") + class DirectionParsing { + + @Test + @DisplayName("Defaults to upstream when null/blank") + void defaultsToUpstream() { + assertEquals(LineagePathBuilder.Direction.UPSTREAM, LineagePathBuilder.Direction.parse(null)); + assertEquals(LineagePathBuilder.Direction.UPSTREAM, LineagePathBuilder.Direction.parse("")); + assertEquals(LineagePathBuilder.Direction.UPSTREAM, LineagePathBuilder.Direction.parse(" ")); + } + + @Test + @DisplayName("Recognized values map case-insensitively") + void recognized() { + assertEquals( + LineagePathBuilder.Direction.UPSTREAM, LineagePathBuilder.Direction.parse("UPSTREAM")); + assertEquals( + LineagePathBuilder.Direction.DOWNSTREAM, + LineagePathBuilder.Direction.parse("Downstream")); + assertEquals(LineagePathBuilder.Direction.BOTH, LineagePathBuilder.Direction.parse(" both ")); + } + + @Test + @DisplayName("Unknown direction is rejected") + void unknown() { + assertThrows( + IllegalArgumentException.class, () -> LineagePathBuilder.Direction.parse("sideways")); + } + } + + @Nested + @DisplayName("Hop budget clamping") + class HopBudgetClamping { + + @Test + @DisplayName("Null and < 1 fall back to default") + void defaults() { + assertEquals(LineagePathBuilder.DEFAULT_MAX_HOPS, LineagePathBuilder.clampMaxHops(null)); + assertEquals(LineagePathBuilder.DEFAULT_MAX_HOPS, LineagePathBuilder.clampMaxHops(0)); + assertEquals(LineagePathBuilder.DEFAULT_MAX_HOPS, LineagePathBuilder.clampMaxHops(-9)); + } + + @Test + @DisplayName("Values within [1, HARD_MAX_HOPS] are passed through") + void passthrough() { + assertEquals(1, LineagePathBuilder.clampMaxHops(1)); + assertEquals(7, LineagePathBuilder.clampMaxHops(7)); + assertEquals( + LineagePathBuilder.HARD_MAX_HOPS, + LineagePathBuilder.clampMaxHops(LineagePathBuilder.HARD_MAX_HOPS)); + } + + @Test + @DisplayName("Values above HARD_MAX_HOPS are clamped down") + void clampedHigh() { + assertEquals(LineagePathBuilder.HARD_MAX_HOPS, LineagePathBuilder.clampMaxHops(9999)); + } + } + + @Nested + @DisplayName("Frontier SPARQL well-formedness") + class FrontierSparql { + + @Test + @DisplayName("Empty frontier is rejected before SPARQL is built") + void emptyFrontier() { + assertThrows( + IllegalArgumentException.class, + () -> LineagePathBuilder.frontierQuery(List.of(), LineagePathBuilder.Direction.UPSTREAM)); + assertThrows( + IllegalArgumentException.class, + () -> LineagePathBuilder.frontierQuery(null, LineagePathBuilder.Direction.UPSTREAM)); + } + + @Test + @DisplayName("Bad URI in the frontier is rejected before SPARQL is built") + void badFrontierUri() { + assertThrows( + IllegalArgumentException.class, + () -> + LineagePathBuilder.frontierQuery( + List.of("not a uri"), LineagePathBuilder.Direction.UPSTREAM)); + } + + @Test + @DisplayName("Upstream query parses and contains both upstream predicates") + void upstreamQueryParses() { + String sparql = + LineagePathBuilder.frontierQuery( + List.of("https://x.com/a"), LineagePathBuilder.Direction.UPSTREAM); + Query q = QueryFactory.create(sparql); + assertTrue(q.isSelectType()); + assertTrue(q.getResultVars().contains("from")); + assertTrue(q.getResultVars().contains("to")); + assertTrue(q.getResultVars().contains("predicate")); + assertTrue(sparql.contains("prov:wasDerivedFrom")); + assertTrue(sparql.contains("om:upstream")); + } + + @Test + @DisplayName("Downstream query inverts prov:wasDerivedFrom") + void downstreamInverts() { + String sparql = + LineagePathBuilder.frontierQuery( + List.of("https://x.com/a"), LineagePathBuilder.Direction.DOWNSTREAM); + QueryFactory.create(sparql); + assertTrue(sparql.contains("?to prov:wasDerivedFrom ?from")); + assertTrue(sparql.contains("om:downstream")); + assertTrue( + sparql.contains("\"^prov:wasDerivedFrom\""), + "Inverted edge must be labelled as ^prov:wasDerivedFrom so callers can render direction"); + } + + @Test + @DisplayName("Both direction emits all four predicate variants") + void bothEmitsAll() { + String sparql = + LineagePathBuilder.frontierQuery( + List.of("https://x.com/a"), LineagePathBuilder.Direction.BOTH); + QueryFactory.create(sparql); + assertTrue(sparql.contains("prov:wasDerivedFrom")); + assertTrue(sparql.contains("om:upstream")); + assertTrue(sparql.contains("om:downstream")); + assertTrue(sparql.contains("^prov:wasDerivedFrom")); + } + + @Test + @DisplayName("Multi-node frontier produces VALUES with all URIs") + void multiNodeValues() { + String sparql = + LineagePathBuilder.frontierQuery( + List.of("https://x.com/a", "https://x.com/b", "https://x.com/c"), + LineagePathBuilder.Direction.UPSTREAM); + QueryFactory.create(sparql); + assertTrue(sparql.contains("")); + assertTrue(sparql.contains("")); + assertTrue(sparql.contains("")); + } + + @Test + @DisplayName("Self-loops are filtered out") + void selfLoopsFiltered() { + String sparql = + LineagePathBuilder.frontierQuery( + List.of("https://x.com/a"), LineagePathBuilder.Direction.UPSTREAM); + assertTrue( + sparql.contains("FILTER(?to != ?from)"), + "Self-loops shouldn't pollute BFS — must be filtered server-side"); + } + } + + @Nested + @DisplayName("Types SPARQL") + class TypesSparql { + + @Test + @DisplayName("Empty input is rejected") + void empty() { + assertThrows(IllegalArgumentException.class, () -> LineagePathBuilder.typesQuery(List.of())); + assertThrows(IllegalArgumentException.class, () -> LineagePathBuilder.typesQuery(null)); + } + + @Test + @DisplayName("Only om: types are returned (filter is present)") + void omFilter() { + String sparql = LineagePathBuilder.typesQuery(List.of("https://x.com/a")); + QueryFactory.create(sparql); + assertTrue(sparql.contains("STRSTARTS"), "Must filter to om-namespaced types"); + assertTrue(sparql.contains("https://open-metadata.org/ontology/")); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LineagePathFinderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LineagePathFinderTest.java new file mode 100644 index 000000000000..a8ae15df2a15 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LineagePathFinderTest.java @@ -0,0 +1,423 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.openmetadata.service.rdf.RdfRepository; + +/** + * Tests for {@link LineagePathFinder}. Two kinds of tests: + * + *
    + *
  1. Pure parsing — {@link LineagePathFinder#parseFrontierResult} and + * {@link LineagePathFinder#parseTypesResult} given hand-built SPARQL JSON. + *
  2. End-to-end BFS — {@link LineagePathFinder#findPath} against a mocked {@link RdfRepository}. + * The mock dispatches on SPARQL shape (frontier vs types) and on which frontier URIs are + * embedded in the query, so each test reads as a tiny graph definition. + *
+ */ +class LineagePathFinderTest { + + private static final String EMPTY = "{\"results\":{\"bindings\":[]}}"; + + private static String frontierRow(String from, String to, String predicate) { + return "{\"from\":{\"value\":\"" + + from + + "\"},\"to\":{\"value\":\"" + + to + + "\"},\"predicate\":{\"value\":\"" + + predicate + + "\"}}"; + } + + private static String frontierBody(String... rows) { + return "{\"results\":{\"bindings\":[" + String.join(",", rows) + "]}}"; + } + + private static String typesRow(String node, String type) { + return "{\"node\":{\"value\":\"" + node + "\"},\"type\":{\"value\":\"" + type + "\"}}"; + } + + private static String typesBody(String... rows) { + return "{\"results\":{\"bindings\":[" + String.join(",", rows) + "]}}"; + } + + @Nested + @DisplayName("Frontier result parsing") + class FrontierParsing { + + @Test + @DisplayName("Empty / null / blank input → empty map") + void empty() { + Set visited = new HashSet<>(); + assertTrue(LineagePathFinder.parseFrontierResult(null, visited).isEmpty()); + assertTrue(LineagePathFinder.parseFrontierResult("", visited).isEmpty()); + assertTrue(LineagePathFinder.parseFrontierResult("not json", visited).isEmpty()); + assertTrue( + LineagePathFinder.parseFrontierResult("{\"results\":{\"bindings\":[]}}", visited) + .isEmpty()); + } + + @Test + @DisplayName("Already-visited 'to' nodes are dropped") + void visitedDropped() { + Set visited = Set.of("urn:b"); + Map next = + LineagePathFinder.parseFrontierResult( + frontierBody(frontierRow("urn:a", "urn:b", "prov:wasDerivedFrom")), visited); + assertTrue(next.isEmpty(), "Already-visited target must not be re-added to next frontier"); + } + + @Test + @DisplayName("First parent wins when multiple frontier rows mention same target") + void firstParentWins() { + Set visited = new HashSet<>(); + Map next = + LineagePathFinder.parseFrontierResult( + frontierBody( + frontierRow("urn:a", "urn:b", "prov:wasDerivedFrom"), + frontierRow("urn:c", "urn:b", "om:upstream")), + visited); + assertEquals("urn:a", next.get("urn:b").parent()); + assertEquals("prov:wasDerivedFrom", next.get("urn:b").predicate()); + } + + @Test + @DisplayName("Rows with missing fields are skipped") + void partialRows() { + String partial = "{\"results\":{\"bindings\":[{\"from\":{\"value\":\"urn:a\"}}]}}"; + assertTrue(LineagePathFinder.parseFrontierResult(partial, new HashSet<>()).isEmpty()); + } + } + + @Nested + @DisplayName("Types result parsing") + class TypesParsing { + + @Test + @DisplayName("Empty / null / blank input → empty map") + void empty() { + assertTrue(LineagePathFinder.parseTypesResult(null).isEmpty()); + assertTrue(LineagePathFinder.parseTypesResult("").isEmpty()); + assertTrue(LineagePathFinder.parseTypesResult("garbage").isEmpty()); + } + + @Test + @DisplayName("Multiple types per node are aggregated and deduplicated") + void multipleTypes() { + String body = + typesBody( + typesRow("urn:t", "https://open-metadata.org/ontology/Table"), + typesRow("urn:t", "https://open-metadata.org/ontology/DataAsset"), + typesRow("urn:t", "https://open-metadata.org/ontology/Table")); + Map> result = LineagePathFinder.parseTypesResult(body); + List types = result.get("urn:t"); + assertEquals(2, types.size(), "Duplicate types must be dropped"); + assertTrue(types.contains("https://open-metadata.org/ontology/Table")); + assertTrue(types.contains("https://open-metadata.org/ontology/DataAsset")); + } + } + + @Nested + @DisplayName("End-to-end BFS") + class EndToEnd { + + private static final String A = "https://open-metadata.org/instance/Table/a"; + private static final String B = "https://open-metadata.org/instance/Table/b"; + private static final String C = "https://open-metadata.org/instance/Table/c"; + private static final String D = "https://open-metadata.org/instance/Table/d"; + + private RdfRepository mockRepo() { + RdfRepository repo = mock(RdfRepository.class); + lenient().when(repo.executeSparqlQuery(anyString(), anyString())).thenReturn(EMPTY); + return repo; + } + + @Test + @DisplayName("from == to: trivial single-node path returned immediately") + void trivialIdentity() { + RdfRepository repo = mockRepo(); + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, A, LineagePathBuilder.Direction.UPSTREAM, 6); + assertTrue(path.found()); + assertEquals(0, path.hops()); + assertEquals(1, path.nodes().size()); + assertEquals(A, path.nodes().get(0).node()); + assertNull(path.nodes().get(0).predicate()); + } + + @Test + @DisplayName("Direct neighbour found in one hop") + void directNeighbour() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) return typesBody(); + if (sparql.contains("<" + A + ">")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, B, LineagePathBuilder.Direction.UPSTREAM, 6); + assertTrue(path.found()); + assertEquals(1, path.hops()); + assertEquals(List.of(A, B), nodeUris(path)); + assertNull(path.nodes().get(0).predicate()); + assertEquals("prov:wasDerivedFrom", path.nodes().get(1).predicate()); + } + + @Test + @DisplayName("Multi-hop A → B → C → D resolves with three predicate-tagged hops") + void multiHop() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) return typesBody(); + if (sparql.contains("<" + A + ">")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + if (sparql.contains("<" + B + ">")) { + return frontierBody(frontierRow(B, C, "om:upstream")); + } + if (sparql.contains("<" + C + ">")) { + return frontierBody(frontierRow(C, D, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, D, LineagePathBuilder.Direction.UPSTREAM, 6); + assertTrue(path.found()); + assertEquals(3, path.hops()); + assertEquals(List.of(A, B, C, D), nodeUris(path)); + assertEquals("prov:wasDerivedFrom", path.nodes().get(1).predicate()); + assertEquals("om:upstream", path.nodes().get(2).predicate()); + assertEquals("prov:wasDerivedFrom", path.nodes().get(3).predicate()); + } + + @Test + @DisplayName("BFS prefers shorter paths even when a longer one is also reachable") + void bfsShortest() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) return typesBody(); + if (sparql.contains("<" + A + ">")) { + return frontierBody( + frontierRow(A, B, "prov:wasDerivedFrom"), + frontierRow(A, D, "prov:wasDerivedFrom")); + } + if (sparql.contains("<" + B + ">") || sparql.contains("<" + D + ">")) { + return frontierBody(frontierRow(B, C, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, D, LineagePathBuilder.Direction.UPSTREAM, 6); + assertTrue(path.found()); + assertEquals(1, path.hops(), "BFS must take the direct A→D edge, not the A→B→D detour"); + } + + @Test + @DisplayName("Cycle A → B → A does not loop forever; resolves to nearest path") + void cycleSafe() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) return typesBody(); + if (sparql.contains("<" + A + ">")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + if (sparql.contains("<" + B + ">")) { + return frontierBody( + frontierRow(B, A, "prov:wasDerivedFrom"), + frontierRow(B, C, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, C, LineagePathBuilder.Direction.UPSTREAM, 6); + assertTrue(path.found()); + assertEquals(2, path.hops()); + assertEquals(List.of(A, B, C), nodeUris(path)); + } + + @Test + @DisplayName("Disconnected target: BFS exhausts the frontier and reports found=false") + void disconnected() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) return typesBody(); + if (sparql.contains("<" + A + ">")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, D, LineagePathBuilder.Direction.UPSTREAM, 6); + assertFalse(path.found()); + assertEquals(0, path.hops()); + assertTrue(path.nodes().isEmpty()); + assertEquals(A, path.from()); + assertEquals(D, path.to()); + } + + @Test + @DisplayName("maxHops budget is honoured; deeper targets are not found") + void maxHopsBudget() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) return typesBody(); + if (sparql.contains("<" + A + ">")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + if (sparql.contains("<" + B + ">")) { + return frontierBody(frontierRow(B, C, "prov:wasDerivedFrom")); + } + if (sparql.contains("<" + C + ">")) { + return frontierBody(frontierRow(C, D, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, D, LineagePathBuilder.Direction.UPSTREAM, 2); + assertFalse(path.found(), "Three hops cannot fit in a budget of two"); + assertEquals(2, path.maxHops()); + } + + @Test + @DisplayName("SPARQL exception during frontier expansion → not-found, no exception bubbles") + void sparqlError() { + RdfRepository repo = mock(RdfRepository.class); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenThrow(new RuntimeException("Fuseki down")); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, B, LineagePathBuilder.Direction.UPSTREAM, 6); + assertFalse(path.found()); + } + + @Test + @DisplayName("Bad URI: validation fires before any SPARQL is sent") + void badUri() { + RdfRepository repo = mock(RdfRepository.class); + LineagePathFinder finder = new LineagePathFinder(repo); + assertThrows( + IllegalArgumentException.class, + () -> finder.findPath("not a uri", B, LineagePathBuilder.Direction.UPSTREAM, 6)); + assertThrows( + IllegalArgumentException.class, + () -> finder.findPath(A, "ftp://x.com/y", LineagePathBuilder.Direction.UPSTREAM, 6)); + } + + @Test + @DisplayName("Type decoration: each path node carries its om: rdf:types") + void typeDecoration() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) { + return typesBody( + typesRow(A, "https://open-metadata.org/ontology/Table"), + typesRow(A, "https://open-metadata.org/ontology/DataAsset"), + typesRow(B, "https://open-metadata.org/ontology/Table")); + } + if (sparql.contains("<" + A + ">") && sparql.contains("?from ?to ?predicate")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, B, LineagePathBuilder.Direction.UPSTREAM, 6); + assertTrue(path.found()); + assertEquals(2, path.nodes().get(0).rdfTypes().size()); + assertEquals(1, path.nodes().get(1).rdfTypes().size()); + } + + @Test + @DisplayName("Type decoration failure does not break the path response") + void typeDecorationFailure() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) { + throw new RuntimeException("Fuseki blip on type query"); + } + if (sparql.contains("<" + A + ">")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + + LineagePathFinder.Path path = + new LineagePathFinder(repo).findPath(A, B, LineagePathBuilder.Direction.UPSTREAM, 6); + assertTrue(path.found(), "Path must still be returned even if type decoration blows up"); + assertEquals(2, path.nodes().size()); + assertNotNull(path.nodes().get(0).rdfTypes()); + assertTrue(path.nodes().get(0).rdfTypes().isEmpty()); + } + + @Test + @DisplayName("Direction defaults to upstream when null is passed") + void directionDefault() { + RdfRepository repo = mockRepo(); + when(repo.executeSparqlQuery(anyString(), anyString())) + .thenAnswer( + inv -> { + String sparql = inv.getArgument(0); + if (sparql.contains("?node ?type")) return typesBody(); + if (sparql.contains("<" + A + ">")) { + return frontierBody(frontierRow(A, B, "prov:wasDerivedFrom")); + } + return EMPTY; + }); + LineagePathFinder.Path path = new LineagePathFinder(repo).findPath(A, B, null, 6); + assertTrue(path.found()); + assertEquals("upstream", path.direction()); + } + } + + private static List nodeUris(LineagePathFinder.Path path) { + return path.nodes().stream().map(LineagePathFinder.Hop::node).toList(); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LouvainTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LouvainTest.java new file mode 100644 index 000000000000..a3c49bca8e73 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/LouvainTest.java @@ -0,0 +1,233 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +class LouvainTest { + + private static Map> g(Object... edges) { + Map> graph = new LinkedHashMap<>(); + for (int i = 0; i < edges.length; i += 3) { + String from = (String) edges[i]; + String to = (String) edges[i + 1]; + double w = ((Number) edges[i + 2]).doubleValue(); + graph.computeIfAbsent(from, k -> new HashMap<>()).put(to, w); + graph.computeIfAbsent(to, k -> new HashMap<>()); + } + return graph; + } + + @Nested + @DisplayName("Constructor / input guards") + class Guards { + + @Test + @DisplayName("maxIterations < 1 is rejected") + void badMaxIterations() { + assertThrows(IllegalArgumentException.class, () -> new Louvain(0)); + assertThrows(IllegalArgumentException.class, () -> new Louvain(-3)); + } + + @Test + @DisplayName("Null and empty graphs return empty result with modularity 0") + void emptyInputs() { + Louvain.Result r1 = new Louvain().compute(null); + assertTrue(r1.communityByNode().isEmpty()); + assertEquals(0.0, r1.modularity()); + Louvain.Result r2 = new Louvain().compute(Map.of()); + assertTrue(r2.communityByNode().isEmpty()); + } + + @Test + @DisplayName("Graph with only self-loops produces singletons (self-loops ignored)") + void onlySelfLoops() { + Map> graph = + new LinkedHashMap<>( + Map.of( + "a", new HashMap<>(Map.of("a", 1.0)), + "b", new HashMap<>(Map.of("b", 1.0)))); + Louvain.Result r = new Louvain().compute(graph); + assertEquals(2, r.communityCount(), "Self-loops carry no community signal"); + } + + @Test + @DisplayName("Negative weights are clamped to zero") + void negativeWeights() { + Louvain.Result r = new Louvain().compute(g("a", "b", -100.0)); + assertEquals(2, r.communityCount(), "Negative-weight edges must not pull nodes together"); + } + } + + @Nested + @DisplayName("Topology → community structure") + class Topology { + + @Test + @DisplayName("Triangle (three nodes, three edges): everyone joins one community") + void triangle() { + Map> graph = g("a", "b", 1.0, "b", "c", 1.0, "a", "c", 1.0); + Louvain.Result r = new Louvain().compute(graph); + assertEquals(1, r.communityCount()); + } + + @Test + @DisplayName("Two cliques connected by a single light edge: two communities") + void twoCliques() { + Map> graph = + g( + "a", "b", 1.0, "b", "c", 1.0, "a", "c", 1.0, "x", "y", 1.0, "y", "z", 1.0, "x", "z", + 1.0, "c", "x", 0.01); + Louvain.Result r = new Louvain().compute(graph); + assertEquals(2, r.communityCount()); + assertEquals(r.communityByNode().get("a"), r.communityByNode().get("b")); + assertEquals(r.communityByNode().get("a"), r.communityByNode().get("c")); + assertEquals(r.communityByNode().get("x"), r.communityByNode().get("y")); + assertEquals(r.communityByNode().get("x"), r.communityByNode().get("z")); + assertNotEquals(r.communityByNode().get("a"), r.communityByNode().get("x")); + } + + @Test + @DisplayName("Disconnected components yield distinct communities") + void disconnectedComponents() { + Map> graph = g("a", "b", 1.0, "c", "d", 1.0, "e", "f", 1.0); + Louvain.Result r = new Louvain().compute(graph); + assertEquals(3, r.communityCount()); + } + + @Test + @DisplayName("Star (hub + leaves) collapses to a single community") + void star() { + Map> graph = new LinkedHashMap<>(); + graph.put("hub", new HashMap<>()); + for (int i = 0; i < 6; i++) { + graph.computeIfAbsent("hub", k -> new HashMap<>()).put("leaf-" + i, 1.0); + graph.put("leaf-" + i, new HashMap<>()); + } + Louvain.Result r = new Louvain().compute(graph); + assertEquals(1, r.communityCount()); + } + + @Test + @DisplayName("Heavy edges pull nodes together against light competing edges") + void edgeWeightsRespected() { + Map> graph = g("a", "b", 100.0, "a", "c", 0.1, "c", "d", 100.0); + Louvain.Result r = new Louvain().compute(graph); + assertEquals(r.communityByNode().get("a"), r.communityByNode().get("b")); + assertEquals(r.communityByNode().get("c"), r.communityByNode().get("d")); + assertNotEquals(r.communityByNode().get("a"), r.communityByNode().get("c")); + } + } + + @Nested + @DisplayName("Symmetrization") + class Symmetrization { + + @Test + @DisplayName("Asymmetric input is treated as undirected") + void asymmetricInput() { + Map> graph = new LinkedHashMap<>(); + graph.put("a", new HashMap<>(Map.of("b", 5.0))); + graph.put("b", new HashMap<>()); + Louvain.Result r = new Louvain().compute(graph); + assertEquals(1, r.communityCount(), "Single edge a→b should still pull a, b together"); + } + + @Test + @DisplayName("Both-directions input doesn't double-influence") + void bothDirectionsSum() { + Map> graph = g("a", "b", 5.0, "b", "a", 5.0); + Louvain.Result r = new Louvain().compute(graph); + assertEquals(1, r.communityCount()); + } + } + + @Nested + @DisplayName("Determinism") + class Determinism { + + @Test + @DisplayName("Repeated runs on the same input produce the same partition") + void deterministic() { + Map> graph = + g( + "a", "b", 1.0, "b", "c", 1.0, "c", "a", 1.0, "x", "y", 1.0, "y", "z", 1.0, "z", "x", + 1.0, "a", "x", 0.05); + Louvain.Result r1 = new Louvain().compute(graph); + Louvain.Result r2 = new Louvain().compute(graph); + assertEquals(r1.communityByNode(), r2.communityByNode()); + assertEquals(r1.modularity(), r2.modularity()); + } + + @Test + @DisplayName("Community ids are dense [0..k-1] in discovery order") + void denseIds() { + Map> graph = g("a", "b", 1.0, "c", "d", 1.0); + Louvain.Result r = new Louvain().compute(graph); + List ids = r.communityByNode().values().stream().distinct().sorted().toList(); + assertEquals(List.of(0, 1), ids); + } + } + + @Nested + @DisplayName("Modularity behaviour") + class Modularity { + + @Test + @DisplayName("Tight clusters produce higher modularity than mixed input") + void clustersHaveHigherQ() { + Map> tight = + g( + "a", "b", 1.0, "b", "c", 1.0, "c", "a", 1.0, "x", "y", 1.0, "y", "z", 1.0, "z", "x", + 1.0, "a", "x", 0.05); + Map> mixed = + g( + "a", "b", 1.0, "a", "c", 1.0, "a", "d", 1.0, "a", "e", 1.0, "a", "f", 1.0, "a", "g", + 1.0); + double qTight = new Louvain().compute(tight).modularity(); + double qMixed = new Louvain().compute(mixed).modularity(); + assertTrue( + qTight > qMixed, + "Two-clique partition must score higher modularity than a star (got tight=" + + qTight + + ", mixed=" + + qMixed + + ")"); + } + } + + @Nested + @DisplayName("Result helpers") + class ResultHelpers { + + @Test + @DisplayName("membersByCommunity is the inverse view of communityByNode") + void membersByCommunity() { + Map> graph = + g( + "a", "b", 1.0, "b", "c", 1.0, "a", "c", 1.0, "x", "y", 1.0, "y", "z", 1.0, "x", "z", + 1.0, "c", "x", 0.01); + Louvain.Result r = new Louvain().compute(graph); + Map> members = r.membersByCommunity(); + assertEquals(2, members.size()); + int totalMembers = members.values().stream().mapToInt(List::size).sum(); + assertEquals(6, totalMembers); + } + + @Test + @DisplayName("Iteration count is non-zero whenever any edge exists") + void iterationsAdvance() { + Louvain.Result r = new Louvain().compute(g("a", "b", 1.0)); + assertTrue(r.iterations() >= 1); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/PageRankTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/PageRankTest.java new file mode 100644 index 000000000000..b6de7085aa87 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/PageRankTest.java @@ -0,0 +1,274 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +/** + * Correctness + edge-case tests for the hand-rolled PageRank implementation. + * + *

Each test names the property under test and uses a small known graph so the expected + * scores can be reasoned about directly. We don't assert exact values from the literature + * (those depend on damping factor and tolerance choices) — instead we assert qualitative + * properties that must hold: + * + *

    + *
  • scores are normalized (sum to 1.0) + *
  • nodes with more incoming weight rank higher + *
  • dangling nodes get a non-zero score (mass redistribution) + *
  • disconnected components both contribute to the result + *
+ */ +class PageRankTest { + + private static Map> g() { + return new LinkedHashMap<>(); + } + + private static void edge(Map> g, String from, String to, double w) { + g.computeIfAbsent(from, k -> new HashMap<>()).put(to, w); + } + + @Nested + @DisplayName("Constructor validation") + class Construction { + + @Test + @DisplayName("Damping outside (0,1) is rejected") + void invalidDamping() { + assertThrows(IllegalArgumentException.class, () -> new PageRank(0.0, 100, 1e-6)); + assertThrows(IllegalArgumentException.class, () -> new PageRank(1.0, 100, 1e-6)); + assertThrows(IllegalArgumentException.class, () -> new PageRank(-0.5, 100, 1e-6)); + } + + @Test + @DisplayName("Non-positive maxIterations is rejected") + void invalidIterations() { + assertThrows(IllegalArgumentException.class, () -> new PageRank(0.85, 0, 1e-6)); + assertThrows(IllegalArgumentException.class, () -> new PageRank(0.85, -1, 1e-6)); + } + + @Test + @DisplayName("Non-positive tolerance is rejected") + void invalidTolerance() { + assertThrows(IllegalArgumentException.class, () -> new PageRank(0.85, 100, 0.0)); + assertThrows(IllegalArgumentException.class, () -> new PageRank(0.85, 100, -1.0)); + } + } + + @Nested + @DisplayName("Edge cases") + class EdgeCases { + + @Test + @DisplayName("Empty graph returns empty result, zero iterations") + void emptyGraph() { + PageRank.Result r = new PageRank().compute(g()); + assertTrue(r.scores().isEmpty()); + assertEquals(0, r.iterations()); + assertTrue(r.converged()); + } + + @Test + @DisplayName("Single node with no edges → score 1.0") + void singleNode() { + Map> g = g(); + g.put("A", new HashMap<>()); + PageRank.Result r = new PageRank().compute(g); + assertEquals(1.0, r.scores().get("A"), 1e-9); + } + + @Test + @DisplayName("Two disconnected nodes get equal score 0.5") + void twoDisconnected() { + Map> g = g(); + g.put("A", new HashMap<>()); + g.put("B", new HashMap<>()); + PageRank.Result r = new PageRank().compute(g); + assertEquals(0.5, r.scores().get("A"), 1e-9); + assertEquals(0.5, r.scores().get("B"), 1e-9); + } + + @Test + @DisplayName("Dangling target node still receives a score") + void danglingTarget() { + Map> g = g(); + edge(g, "A", "B", 1.0); + PageRank.Result r = new PageRank().compute(g); + assertTrue(r.scores().containsKey("A")); + assertTrue(r.scores().containsKey("B")); + assertTrue(r.scores().get("B") > 0); + assertTrue( + r.scores().get("B") > r.scores().get("A"), + "B has incoming edge from A so should score higher than A"); + } + + @Test + @DisplayName("Self-loop on a single node → still normalized") + void selfLoop() { + Map> g = g(); + edge(g, "A", "A", 1.0); + PageRank.Result r = new PageRank().compute(g); + assertEquals(1.0, r.scores().get("A"), 1e-9); + } + + @Test + @DisplayName("Edge with zero weight contributes nothing") + void zeroWeightEdgeIgnored() { + Map> g = g(); + edge(g, "A", "B", 0.0); + // A is effectively dangling; A and B should split mass via dangling redistribution. + PageRank.Result r = new PageRank().compute(g); + assertEquals(0.5, r.scores().get("A"), 1e-3); + assertEquals(0.5, r.scores().get("B"), 1e-3); + } + } + + @Nested + @DisplayName("Output normalization") + class Normalization { + + @Test + @DisplayName("Scores sum to 1.0 across the graph") + void scoresSumToOne() { + Map> g = g(); + edge(g, "A", "B", 1.0); + edge(g, "B", "C", 1.0); + edge(g, "C", "A", 1.0); + edge(g, "C", "B", 1.0); + PageRank.Result r = new PageRank().compute(g); + double total = 0; + for (double v : r.scores().values()) total += v; + assertEquals(1.0, total, 1e-6); + } + + @Test + @DisplayName("Symmetric graph produces equal scores") + void symmetricGraph() { + Map> g = g(); + edge(g, "A", "B", 1.0); + edge(g, "B", "A", 1.0); + PageRank.Result r = new PageRank().compute(g); + assertEquals(r.scores().get("A"), r.scores().get("B"), 1e-6); + } + } + + @Nested + @DisplayName("Ranking properties") + class RankingProperties { + + @Test + @DisplayName("Hub node (many incoming edges) ranks highest") + void hubRanksHighest() { + Map> g = g(); + // A, B, C, D all point at HUB. HUB has no outgoing edges. + edge(g, "A", "HUB", 1.0); + edge(g, "B", "HUB", 1.0); + edge(g, "C", "HUB", 1.0); + edge(g, "D", "HUB", 1.0); + PageRank.Result r = new PageRank().compute(g); + double hub = r.scores().get("HUB"); + for (String n : new String[] {"A", "B", "C", "D"}) { + assertTrue(hub > r.scores().get(n), "HUB > " + n + ": " + r.scores()); + } + } + + @Test + @DisplayName("Edge weight matters: heavy-weighted target outranks lightly-weighted target") + void edgeWeightMatters() { + Map> g = g(); + // SOURCE → HEAVY (weight 10), SOURCE → LIGHT (weight 0.1) + edge(g, "SOURCE", "HEAVY", 10.0); + edge(g, "SOURCE", "LIGHT", 0.1); + PageRank.Result r = new PageRank().compute(g); + assertTrue( + r.scores().get("HEAVY") > r.scores().get("LIGHT"), + "Heavy edge should outrank light edge: " + r.scores()); + } + + @Test + @DisplayName("Star topology — center outranks every leaf") + void starTopology() { + Map> g = g(); + for (int i = 0; i < 10; i++) { + edge(g, "leaf-" + i, "center", 1.0); + } + PageRank.Result r = new PageRank().compute(g); + double center = r.scores().get("center"); + for (int i = 0; i < 10; i++) { + assertTrue(center > r.scores().get("leaf-" + i)); + } + } + + @Test + @DisplayName("Two-component graph: each component ranks consistently within itself") + void twoComponents() { + Map> g = g(); + // Component 1: A → B → C (chain) + edge(g, "A", "B", 1.0); + edge(g, "B", "C", 1.0); + // Component 2: X → Y → X (cycle) + edge(g, "X", "Y", 1.0); + edge(g, "Y", "X", 1.0); + PageRank.Result r = new PageRank().compute(g); + assertTrue(r.scores().get("X") > 0); + assertTrue(r.scores().get("Y") > 0); + assertTrue(r.scores().get("C") > 0, "Dangling end of chain still gets mass"); + // Within the cycle, X and Y should be equal + assertEquals(r.scores().get("X"), r.scores().get("Y"), 1e-6); + } + } + + @Nested + @DisplayName("Convergence") + class Convergence { + + @Test + @DisplayName("Small graphs converge in well under maxIterations") + void convergesQuickly() { + Map> g = g(); + edge(g, "A", "B", 1.0); + edge(g, "B", "A", 1.0); + PageRank.Result r = new PageRank().compute(g); + assertTrue(r.converged()); + assertTrue(r.iterations() < 50, "Took too many iterations: " + r.iterations()); + } + + @Test + @DisplayName("Tight tolerance still converges") + void tightTolerance() { + Map> g = g(); + edge(g, "A", "B", 1.0); + edge(g, "B", "C", 1.0); + edge(g, "C", "A", 1.0); + PageRank.Result r = new PageRank(0.85, 1000, 1e-12).compute(g); + assertTrue(r.converged()); + } + + @Test + @DisplayName("maxIterations=1 returns without converging on a non-trivial graph") + void maxIterationsHonored() { + Map> g = g(); + // Hub-and-spoke needs more than 1 iteration to settle. + for (int i = 0; i < 10; i++) edge(g, "leaf-" + i, "center", 1.0); + PageRank.Result r = new PageRank(0.85, 1, 1e-12).compute(g); + assertEquals(1, r.iterations()); + } + } + + @Test + @DisplayName("nodes() helper returns the union of sources and targets") + void nodesHelper() { + Map> g = g(); + edge(g, "A", "B", 1.0); + edge(g, "B", "C", 1.0); + assertEquals(java.util.Set.of("A", "B", "C"), PageRank.nodes(g)); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/RecommendationsQueryBuilderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/RecommendationsQueryBuilderTest.java new file mode 100644 index 000000000000..9204d00fe957 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/insights/RecommendationsQueryBuilderTest.java @@ -0,0 +1,108 @@ +package org.openmetadata.service.rdf.insights; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.apache.jena.query.Query; +import org.apache.jena.query.QueryFactory; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +class RecommendationsQueryBuilderTest { + + private static final String URI = "https://open-metadata.org/instance/Table/abc"; + + @Nested + @DisplayName("Input validation") + class InputValidation { + + @Test + @DisplayName("Bad URI is rejected before SPARQL is generated") + void badUri() { + assertThrows( + IllegalArgumentException.class, () -> RecommendationsQueryBuilder.build(null, 10)); + assertThrows( + IllegalArgumentException.class, () -> RecommendationsQueryBuilder.build(" ", 10)); + assertThrows( + IllegalArgumentException.class, + () -> RecommendationsQueryBuilder.build("ftp://x.com/y", 10)); + assertThrows( + IllegalArgumentException.class, + () -> RecommendationsQueryBuilder.build("http://x.com/y> ; DROP", 10)); + } + + @Test + @DisplayName("Limit below 1 is clamped to 1") + void limitClampedLow() { + String sparql = RecommendationsQueryBuilder.build(URI, -10); + assertTrue(sparql.endsWith("LIMIT 1")); + } + + @Test + @DisplayName("Limit above MAX_LIMIT is clamped down") + void limitClampedHigh() { + String sparql = RecommendationsQueryBuilder.build(URI, 9999); + assertTrue(sparql.endsWith("LIMIT " + RecommendationsQueryBuilder.MAX_LIMIT)); + } + } + + @Nested + @DisplayName("SPARQL well-formedness") + class WellFormed { + + @Test + @DisplayName("Generated query parses with Jena and selects the expected vars") + void parses() { + Query q = QueryFactory.create(RecommendationsQueryBuilder.build(URI, 10)); + assertTrue(q.isSelectType()); + assertTrue(q.getResultVars().contains("candidate")); + assertTrue(q.getResultVars().contains("tagOverlap")); + assertTrue(q.getResultVars().contains("glossaryOverlap")); + assertTrue(q.getResultVars().contains("lineageOverlap")); + assertTrue(q.getResultVars().contains("score")); + } + + @Test + @DisplayName("Lineage neighbour predicates cover both directions and the prov inverse") + void lineagePredicates() { + String sparql = RecommendationsQueryBuilder.build(URI, 10); + assertTrue(sparql.contains("om:upstream")); + assertTrue(sparql.contains("om:downstream")); + assertTrue(sparql.contains("prov:wasDerivedFrom")); + assertTrue(sparql.contains("^prov:wasDerivedFrom")); + } + + @Test + @DisplayName("Score formula uses the documented weights and adds three terms") + void scoreFormula() { + String sparql = RecommendationsQueryBuilder.build(URI, 10); + assertTrue(sparql.contains(Double.toString(RecommendationsQueryBuilder.WEIGHT_TAG))); + assertTrue(sparql.contains(Double.toString(RecommendationsQueryBuilder.WEIGHT_GLOSSARY))); + assertTrue(sparql.contains(Double.toString(RecommendationsQueryBuilder.WEIGHT_LINEAGE))); + assertTrue(sparql.contains("ORDER BY DESC(?score)")); + } + + @Test + @DisplayName("Each sub-SELECT excludes the seed itself") + void excludesSeed() { + String sparql = RecommendationsQueryBuilder.build(URI, 10); + long filterCount = + sparql + .lines() + .filter(line -> line.contains("FILTER(?candidate != <" + URI + ">)")) + .count(); + assertTrue(filterCount >= 3, "All three sub-SELECTs must filter out the seed itself"); + } + + @Test + @DisplayName("Outer GROUP BY/SUM combines per-dimension partial counts") + void groupBySum() { + String sparql = RecommendationsQueryBuilder.build(URI, 10); + assertTrue(sparql.contains("SUM(?t)")); + assertTrue(sparql.contains("SUM(?g)")); + assertTrue(sparql.contains("SUM(?l)")); + assertTrue(sparql.contains("GROUP BY ?candidate")); + } + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/rdf/translator/RdfUsageMapperTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/translator/RdfUsageMapperTest.java new file mode 100644 index 000000000000..d33b95881a2c --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/rdf/translator/RdfUsageMapperTest.java @@ -0,0 +1,144 @@ +package org.openmetadata.service.rdf.translator; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.ModelFactory; +import org.apache.jena.rdf.model.Property; +import org.apache.jena.rdf.model.Resource; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +class RdfUsageMapperTest { + + private static final String OM = "https://open-metadata.org/ontology/"; + + private ObjectMapper mapper; + private Model model; + private Resource entity; + + @BeforeEach + void setUp() { + mapper = new ObjectMapper(); + model = ModelFactory.createDefaultModel(); + entity = model.createResource("https://open-metadata.org/entity/table/abc"); + } + + @Test + @DisplayName("Full usage summary emits count + percentile triples for daily/weekly/monthly") + void fullUsageSummary() { + ObjectNode usage = mapper.createObjectNode(); + putStats(usage, "dailyStats", 1234, 92.5); + putStats(usage, "weeklyStats", 8500, 88.0); + putStats(usage, "monthlyStats", 35_000, 90.1); + usage.put("date", "2026-04-29"); + + RdfUsageMapper.emitUsageSummary(usage, entity, model); + + assertCount("usageDailyCount", 1234); + assertCount("usageWeeklyCount", 8500); + assertCount("usageMonthlyCount", 35_000); + assertPercentile("usageDailyPercentile", 92.5); + assertPercentile("usageWeeklyPercentile", 88.0); + assertPercentile("usageMonthlyPercentile", 90.1); + + Property usageDate = model.createProperty(OM, "usageDate"); + assertTrue(model.contains(entity, usageDate)); + assertEquals("2026-04-29", model.getProperty(entity, usageDate).getString()); + } + + @Test + @DisplayName("Null usageSummary is a no-op") + void nullUsage() { + RdfUsageMapper.emitUsageSummary(null, entity, model); + assertEquals(0, model.size()); + } + + @Test + @DisplayName("Non-object usageSummary (string, array) is a no-op") + void nonObjectUsage() { + RdfUsageMapper.emitUsageSummary(mapper.getNodeFactory().textNode("oops"), entity, model); + RdfUsageMapper.emitUsageSummary(mapper.createArrayNode(), entity, model); + assertEquals(0, model.size()); + } + + @Test + @DisplayName("Missing percentileRank is allowed — count is still emitted") + void countWithoutPercentile() { + ObjectNode usage = mapper.createObjectNode(); + ObjectNode dailyStats = mapper.createObjectNode(); + dailyStats.put("count", 42); + usage.set("dailyStats", dailyStats); + + RdfUsageMapper.emitUsageSummary(usage, entity, model); + + assertCount("usageDailyCount", 42); + assertFalse( + model.contains(entity, model.createProperty(OM, "usageDailyPercentile")), + "Percentile must not be emitted when not present"); + } + + @Test + @DisplayName("Non-numeric count or percentile is silently skipped") + void nonNumericValuesSkipped() { + ObjectNode usage = mapper.createObjectNode(); + ObjectNode bad = mapper.createObjectNode(); + bad.put("count", "not-a-number"); + bad.put("percentileRank", "very high"); + usage.set("dailyStats", bad); + + RdfUsageMapper.emitUsageSummary(usage, entity, model); + assertEquals(0, model.size(), "Non-numeric stats must be ignored, not coerced"); + } + + @Test + @DisplayName("Only weekly stats present — daily / monthly predicates absent") + void onlyWeekly() { + ObjectNode usage = mapper.createObjectNode(); + putStats(usage, "weeklyStats", 100, 50.0); + + RdfUsageMapper.emitUsageSummary(usage, entity, model); + + assertFalse(model.contains(entity, model.createProperty(OM, "usageDailyCount"))); + assertCount("usageWeeklyCount", 100); + assertFalse(model.contains(entity, model.createProperty(OM, "usageMonthlyCount"))); + } + + @Test + @DisplayName("Date is emitted as xsd:date typed literal") + void datePresent() { + ObjectNode usage = mapper.createObjectNode(); + usage.put("date", "2026-04-29"); + RdfUsageMapper.emitUsageSummary(usage, entity, model); + + Property usageDate = model.createProperty(OM, "usageDate"); + assertTrue( + model.contains(entity, usageDate), + "Date should be present even when no stats are recorded"); + assertEquals("2026-04-29", model.getProperty(entity, usageDate).getString()); + } + + private void putStats(ObjectNode usage, String key, long count, double percentile) { + ObjectNode stats = mapper.createObjectNode(); + stats.put("count", count); + stats.put("percentileRank", percentile); + usage.set(key, stats); + } + + private void assertCount(String predicate, long expected) { + Property p = model.createProperty(OM, predicate); + assertTrue(model.contains(entity, p), "Expected predicate " + predicate); + assertEquals(expected, model.getProperty(entity, p).getLong()); + } + + private void assertPercentile(String predicate, double expected) { + Property p = model.createProperty(OM, predicate); + assertTrue(model.contains(entity, p), "Expected predicate " + predicate); + assertEquals(expected, model.getProperty(entity, p).getDouble(), 1e-9); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/mcp/McpUsageResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/mcp/McpUsageResourceTest.java deleted file mode 100644 index 29f817ba312b..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/mcp/McpUsageResourceTest.java +++ /dev/null @@ -1,278 +0,0 @@ -/* - * Copyright 2025 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.openmetadata.service.resources.mcp; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.anyLong; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -import jakarta.ws.rs.core.Response; -import jakarta.ws.rs.core.SecurityContext; -import java.security.Principal; -import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Map; -import java.util.UUID; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.MockedConstruction; -import org.mockito.MockedStatic; -import org.mockito.Mockito; -import org.openmetadata.schema.entity.app.App; -import org.openmetadata.schema.entity.app.AppExtension; -import org.openmetadata.schema.entity.app.mcp.McpToolCallUsage; -import org.openmetadata.service.apps.AbstractNativeApplication; -import org.openmetadata.service.apps.ApplicationContext; -import org.openmetadata.service.apps.bundles.mcp.McpAppConstants; -import org.openmetadata.service.jdbi3.AppRepository; -import org.openmetadata.service.security.AuthorizationException; -import org.openmetadata.service.security.Authorizer; - -class McpUsageResourceTest { - - private Authorizer authorizer; - private McpUsageResource resource; - private MockedConstruction appRepositoryConstruction; - private MockedStatic appContextStatic; - private SecurityContext adminContext; - private SecurityContext userContext; - - @BeforeEach - void setUp() { - authorizer = mock(Authorizer.class); - appRepositoryConstruction = - Mockito.mockConstruction(AppRepository.class, (mock, ctx) -> stubRepo(mock)); - - ApplicationContext appContext = mock(ApplicationContext.class); - AbstractNativeApplication nativeApp = mock(AbstractNativeApplication.class); - when(nativeApp.getApp()) - .thenReturn(new App().withId(UUID.randomUUID()).withName(McpAppConstants.MCP_APP_NAME)); - when(appContext.getAppIfExists(McpAppConstants.MCP_APP_NAME)).thenReturn(nativeApp); - appContextStatic = Mockito.mockStatic(ApplicationContext.class); - appContextStatic.when(ApplicationContext::getInstance).thenReturn(appContext); - - resource = new McpUsageResource(authorizer); - adminContext = stubSecurityContext("admin"); - userContext = stubSecurityContext("alice"); - } - - @AfterEach - void tearDown() { - appRepositoryConstruction.close(); - appContextStatic.close(); - } - - @Test - void summaryAggregatesCountsAndExcludesBotsFromUniqueUsers() { - stubRows( - row("search_metadata", "alice", true, daysAgo(2)), - row("search_metadata", "alice", true, daysAgo(1)), - row("create_glossary", "bob", false, daysAgo(1)), - row("search_metadata", "McpApplicationBot", true, daysAgo(1)), - row("search_metadata", "ingestion-bot", true, daysAgo(1))); - - Response response = resource.getSummary(adminContext, null, null); - - Map body = bodyAsMap(response); - assertThat(body.get("total")).isEqualTo(5L); - assertThat(body.get("totalSuccess")).isEqualTo(4L); - assertThat(body.get("totalFailed")).isEqualTo(1L); - assertThat(body.get("uniqueUsers")).isEqualTo(2); - } - - @Test - void summaryDeniesNonAdmin() { - doThrow(new AuthorizationException("forbidden")).when(authorizer).authorizeAdmin(userContext); - try { - resource.getSummary(userContext, null, null); - } catch (AuthorizationException expected) { - return; - } - throw new AssertionError("expected AuthorizationException"); - } - - @Test - void breakdownByToolMatchesSummaryTotal() { - stubRows( - row("search_metadata", "alice", true, daysAgo(1)), - row("search_metadata", "bob", true, daysAgo(1)), - row("create_glossary", "alice", true, daysAgo(1))); - - Response summaryResp = resource.getSummary(adminContext, null, null); - Response toolsResp = resource.getByTool(adminContext, null, null); - - long summaryTotal = ((Number) bodyAsMap(summaryResp).get("total")).longValue(); - long toolsTotal = - ((Map) toolsResp.getEntity()) - .values().stream().mapToLong(Long::longValue).sum(); - assertThat(toolsTotal).isEqualTo(summaryTotal); - } - - @Test - void breakdownByUserExcludesBots() { - stubRows( - row("search_metadata", "alice", true, daysAgo(1)), - row("search_metadata", "McpApplicationBot", true, daysAgo(1)), - row("search_metadata", "SystemBot", true, daysAgo(1)), - row("search_metadata", "ingestion-bot", true, daysAgo(1)), - row("search_metadata", "profiler-bot", true, daysAgo(1)), - row("search_metadata", "robot-overlord", true, daysAgo(1))); - - Response response = resource.getByUser(adminContext, null, null); - - Map body = (Map) response.getEntity(); - assertThat(body).containsOnlyKeys("alice", "robot-overlord"); - assertThat(body.get("alice")).isEqualTo(1L); - assertThat(body.get("robot-overlord")).isEqualTo(1L); - } - - @Test - void historyBucketsByUtcDayAndFillsEmptyDays() { - long today = McpUsageResource.startOfDay(Instant.now().toEpochMilli()); - long yesterday = today - Duration.ofDays(1).toMillis(); - long twoDaysAgo = today - Duration.ofDays(2).toMillis(); - stubRows( - row("a", "alice", true, today + 1000), - row("a", "alice", true, today + 2000), - row("a", "alice", true, twoDaysAgo + 500)); - - Response response = - resource.getHistory(adminContext, twoDaysAgo, today + Duration.ofDays(1).toMillis()); - - Map body = (Map) response.getEntity(); - assertThat(body.get(twoDaysAgo)).isEqualTo(1L); - assertThat(body.get(yesterday)).isEqualTo(0L); - assertThat(body.get(today)).isEqualTo(2L); - } - - @Test - void meReturnsOnlyCallerRows() { - stubRows( - row("search_metadata", "alice", true, daysAgo(1)), - row("search_metadata", "alice", false, daysAgo(1)), - row("create_glossary", "bob", true, daysAgo(1))); - - Response response = resource.getMine(userContext, null, null); - - Map body = bodyAsMap(response); - assertThat(body.get("total")).isEqualTo(2L); - Map byTool = (Map) body.get("byTool"); - assertThat(byTool).containsEntry("search_metadata", 2L); - assertThat(byTool).doesNotContainKey("create_glossary"); - } - - @Test - void invalidWindowReturnsBadRequest() { - long now = Instant.now().toEpochMilli(); - Response response = resource.getSummary(adminContext, now, now - 1); - - assertThat(response.getStatus()).isEqualTo(400); - Map body = bodyAsMap(response); - assertThat(body.get("error")).isEqualTo("startTs must be before endTs"); - } - - @Test - void equalStartAndEndAlsoRejected() { - long now = Instant.now().toEpochMilli(); - Response response = resource.getByTool(adminContext, now, now); - - assertThat(response.getStatus()).isEqualTo(400); - } - - @Test - void mcpAppNotInitializedReturnsZeroCounts() { - appContextStatic.close(); - ApplicationContext emptyContext = mock(ApplicationContext.class); - when(emptyContext.getAppIfExists(McpAppConstants.MCP_APP_NAME)).thenReturn(null); - appContextStatic = Mockito.mockStatic(ApplicationContext.class); - appContextStatic.when(ApplicationContext::getInstance).thenReturn(emptyContext); - - Response response = resource.getSummary(adminContext, null, null); - - assertThat(bodyAsMap(response).get("total")).isEqualTo(0L); - } - - private void stubRepo(AppRepository mock) { - lenient() - .when( - mock.listAppExtensionInWindowByName( - any(App.class), - anyLong(), - anyLong(), - anyInt(), - anyInt(), - eq(McpToolCallUsage.class), - eq(AppExtension.ExtensionType.LIMITS))) - .thenReturn(new ArrayList<>()); - } - - private void stubRows(McpToolCallUsage... rows) { - AppRepository repo = appRepositoryConstruction.constructed().getFirst(); - when(repo.listAppExtensionInWindowByName( - any(App.class), - anyLong(), - anyLong(), - anyInt(), - eq(0), - eq(McpToolCallUsage.class), - eq(AppExtension.ExtensionType.LIMITS))) - .thenReturn(new ArrayList<>(Arrays.asList(rows))); - when(repo.listAppExtensionInWindowByName( - any(App.class), - anyLong(), - anyLong(), - anyInt(), - eq(rows.length), - eq(McpToolCallUsage.class), - eq(AppExtension.ExtensionType.LIMITS))) - .thenReturn(new ArrayList<>()); - } - - private static McpToolCallUsage row(String tool, String user, boolean success, long ts) { - return new McpToolCallUsage() - .withAppId(UUID.randomUUID()) - .withAppName(McpAppConstants.MCP_APP_NAME) - .withExtension(AppExtension.ExtensionType.LIMITS) - .withToolName(tool) - .withUserName(user) - .withSuccess(success) - .withTimestamp(ts); - } - - private static long daysAgo(int days) { - return Instant.now().minus(Duration.ofDays(days)).toEpochMilli(); - } - - private static SecurityContext stubSecurityContext(String name) { - SecurityContext ctx = mock(SecurityContext.class); - Principal principal = mock(Principal.class); - when(principal.getName()).thenReturn(name); - when(ctx.getUserPrincipal()).thenReturn(principal); - return ctx; - } - - @SuppressWarnings("unchecked") - private static Map bodyAsMap(Response response) { - return (Map) response.getEntity(); - } -} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/OntologyDocumentTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/OntologyDocumentTest.java new file mode 100644 index 000000000000..e6c0483e39ec --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/OntologyDocumentTest.java @@ -0,0 +1,70 @@ +package org.openmetadata.service.resources.rdf; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import jakarta.ws.rs.core.Response; +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.ModelFactory; +import org.apache.jena.riot.Lang; +import org.apache.jena.riot.RDFDataMgr; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +class OntologyDocumentTest { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String OWL_VERSION_INFO = "http://www.w3.org/2002/07/owl#versionInfo"; + + @Test + @DisplayName( + "Ontology endpoint serves Turtle by default with the bumped version and core classes") + void testServeTurtle() { + Response response = OntologyDocument.serve("turtle"); + assertEquals(200, response.getStatus()); + String body = response.getEntity().toString(); + assertNotNull(body); + + Model parsed = ModelFactory.createDefaultModel(); + RDFDataMgr.read( + parsed, + new java.io.ByteArrayInputStream(body.getBytes(java.nio.charset.StandardCharsets.UTF_8)), + Lang.TURTLE); + + assertTrue( + parsed.contains( + parsed.createResource(OM_NS), + parsed.createProperty(OWL_VERSION_INFO), + parsed.createLiteral("1.1.0")), + "Ontology document should declare owl:versionInfo \"1.1.0\" on the om: ontology"); + + assertTrue( + parsed.containsResource(parsed.createResource(OM_NS + "Column")), + "Core om:Column class must be present in the served ontology"); + assertTrue( + parsed.containsResource(parsed.createResource(OM_NS + "TableConstraint")), + "Newly added om:TableConstraint class must be present"); + assertTrue( + parsed.containsResource(parsed.createResource(OM_NS + "LineageDetails")), + "Newly declared om:LineageDetails class must be present"); + } + + @Test + @DisplayName("Ontology endpoint can render the same document as JSON-LD") + void testServeJsonLd() { + Response response = OntologyDocument.serve("jsonld"); + assertEquals(200, response.getStatus()); + assertEquals("application/ld+json", response.getMediaType().toString()); + String body = response.getEntity().toString(); + assertTrue(body.contains("@context") || body.contains("@graph")); + } + + @Test + @DisplayName("Unknown format defaults to Turtle") + void testUnknownFormatFallsBackToTurtle() { + Response response = OntologyDocument.serve("nonsense"); + assertEquals(200, response.getStatus()); + assertEquals("text/turtle", response.getMediaType().toString()); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/RdfShaclValidatorTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/RdfShaclValidatorTest.java new file mode 100644 index 000000000000..a97d5167693f --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/rdf/RdfShaclValidatorTest.java @@ -0,0 +1,105 @@ +package org.openmetadata.service.resources.rdf; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.apache.jena.rdf.model.Model; +import org.apache.jena.rdf.model.ModelFactory; +import org.apache.jena.rdf.model.Resource; +import org.apache.jena.shacl.ValidationReport; +import org.apache.jena.vocabulary.RDF; +import org.apache.jena.vocabulary.RDFS; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +class RdfShaclValidatorTest { + + private static final String OM_NS = "https://open-metadata.org/ontology/"; + private static final String BASE = "https://open-metadata.org/"; + + @Test + @DisplayName("A column-lineage edge whose om:fromColumn is a string literal violates the shape") + void testColumnLineageRejectsLiteralFromColumn() { + Model model = ModelFactory.createDefaultModel(); + Resource lineage = model.createResource(BASE + "lineageDetails/x/y/colLineage/1"); + lineage.addProperty(RDF.type, model.createResource(OM_NS + "ColumnLineage")); + // Wrong: literal where the shape requires om:Column. + lineage.addProperty(model.createProperty(OM_NS, "fromColumn"), "service.db.s.t.col_a"); + lineage.addProperty( + model.createProperty(OM_NS, "toColumn"), + model.createResource(BASE + "entity/column/service.db.s.target.col_b")); + + ValidationReport report = RdfShaclValidator.validate(model); + assertFalse( + report.conforms(), + "Literal om:fromColumn should violate ColumnLineageShape (om:Column class constraint)"); + } + + @Test + @DisplayName("Properly-shaped column-lineage with URI references conforms") + void testColumnLineageAcceptsUriReferences() { + Model model = ModelFactory.createDefaultModel(); + Resource fromCol = model.createResource(BASE + "entity/column/service.db.s.t.col_a"); + fromCol.addProperty(RDF.type, model.createResource(OM_NS + "Column")); + fromCol.addProperty(model.createProperty(OM_NS, "fullyQualifiedName"), "service.db.s.t.col_a"); + fromCol.addProperty(RDFS.label, "col_a"); + + Resource toCol = model.createResource(BASE + "entity/column/service.db.s.target.col_b"); + toCol.addProperty(RDF.type, model.createResource(OM_NS + "Column")); + toCol.addProperty( + model.createProperty(OM_NS, "fullyQualifiedName"), "service.db.s.target.col_b"); + toCol.addProperty(RDFS.label, "col_b"); + + Resource lineage = model.createResource(BASE + "lineageDetails/x/y/colLineage/1"); + lineage.addProperty(RDF.type, model.createResource(OM_NS + "ColumnLineage")); + lineage.addProperty(model.createProperty(OM_NS, "fromColumn"), fromCol); + lineage.addProperty(model.createProperty(OM_NS, "toColumn"), toCol); + lineage.addProperty(model.createProperty(OM_NS, "fromColumnFqn"), "service.db.s.t.col_a"); + lineage.addProperty(model.createProperty(OM_NS, "toColumnFqn"), "service.db.s.target.col_b"); + + ValidationReport report = RdfShaclValidator.validate(model); + assertTrue( + report.conforms(), + "URI-based column lineage with both endpoints typed as om:Column should conform: " + + reportSummary(report)); + } + + @Test + @DisplayName("A TableConstraint missing constraintType violates TableConstraintShape") + void testTableConstraintRequiresType() { + Model model = ModelFactory.createDefaultModel(); + Resource constraint = model.createResource(BASE + "entity/table/t/constraint/0"); + constraint.addProperty(RDF.type, model.createResource(OM_NS + "TableConstraint")); + Resource col = model.createResource(BASE + "entity/column/service.db.s.t.id"); + col.addProperty(RDF.type, model.createResource(OM_NS + "Column")); + col.addProperty(model.createProperty(OM_NS, "fullyQualifiedName"), "service.db.s.t.id"); + col.addProperty(RDFS.label, "id"); + constraint.addProperty(model.createProperty(OM_NS, "hasConstrainedColumn"), col); + + ValidationReport report = RdfShaclValidator.validate(model); + assertFalse( + report.conforms(), + "TableConstraint without om:constraintType should violate TableConstraintShape minCount=1"); + } + + @Test + @DisplayName("GlossaryTerm without skos:inScheme violates GlossaryTermShape") + void testGlossaryTermRequiresInScheme() { + Model model = ModelFactory.createDefaultModel(); + Resource term = model.createResource(BASE + "entity/glossaryTerm/123"); + term.addProperty(RDF.type, model.createResource(OM_NS + "GlossaryTerm")); + term.addProperty(RDFS.label, "Customer"); + term.addProperty(model.createProperty(OM_NS, "fullyQualifiedName"), "BusinessTerms.Customer"); + + ValidationReport report = RdfShaclValidator.validate(model); + assertFalse( + report.conforms(), + "GlossaryTerm missing skos:inScheme must be flagged so we surface broken glossary memberships"); + } + + private static String reportSummary(ValidationReport report) { + StringBuilder sb = new StringBuilder(); + report.getEntries().forEach(e -> sb.append(e).append("\n")); + return sb.toString(); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/DbTuneReportTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/DbTuneReportTest.java deleted file mode 100644 index f3e4a254dbd1..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/DbTuneReportTest.java +++ /dev/null @@ -1,173 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.List; -import java.util.Map; -import org.junit.jupiter.api.Test; - -class DbTuneReportTest { - - @Test - void formatBytes_handlesAllScales() { - assertEquals("0 B", DbTuneReport.formatBytes(0)); - assertEquals("512 B", DbTuneReport.formatBytes(512)); - assertEquals("2 KB", DbTuneReport.formatBytes(2048)); - assertEquals("4 MB", DbTuneReport.formatBytes(4L * 1024 * 1024)); - assertEquals("1.5 GB", DbTuneReport.formatBytes((long) (1.5 * 1024 * 1024 * 1024))); - } - - @Test - void formatSettings_emptyOrNullShowsDefault() { - assertEquals("(default)", DbTuneReport.formatSettings(null)); - assertEquals("(default)", DbTuneReport.formatSettings(Map.of())); - } - - @Test - void formatSettings_sortsKeysAlphabetically() { - String formatted = - DbTuneReport.formatSettings( - Map.of( - "autovacuum_vacuum_scale_factor", "0.02", - "autovacuum_analyze_scale_factor", "0.01")); - - assertEquals( - "autovacuum_analyze_scale_factor=0.01, autovacuum_vacuum_scale_factor=0.02", formatted); - } - - @Test - void render_includesEngineAndAllSections() { - DbTuneResult result = - new DbTuneResult( - "PostgreSQL", - "17.2", - List.of( - new ServerParamCheck( - "shared_buffers", "16384", "40% of RAM", ServerParamCheck.STATUS_UNTUNED, "")), - List.of( - new TableRecommendation( - "storage_container_entity", - Action.APPLY, - 580_000, - 2L * 1024 * 1024 * 1024, - Map.of(), - Map.of("autovacuum_vacuum_scale_factor", "0.02"), - "Large entity table"))); - - String report = DbTuneReport.render(result); - - assertTrue(report.contains("PostgreSQL 17.2")); - assertTrue(report.contains("Server-level parameter compliance")); - assertTrue(report.contains("Per-table recommendations")); - assertTrue(report.contains("storage_container_entity")); - assertTrue(report.contains("APPLY")); - assertTrue(report.contains("Next steps:")); - } - - @Test - void render_zeroRecommendationsSuppressesAllMatchAndNextSteps() { - DbTuneResult result = new DbTuneResult("PostgreSQL", "17.2", List.of(), List.of()); - - String report = DbTuneReport.render(result); - - assertTrue(report.contains("none of the tracked tables exist")); - assertFalse( - report.contains("already match their recommended settings"), - "Empty recommendations must not claim everything matches"); - assertFalse(report.contains("Next steps:")); - } - - @Test - void render_noActionableShowsAllGoodMessage() { - DbTuneResult result = - new DbTuneResult( - "PostgreSQL", - "17.2", - List.of(), - List.of( - new TableRecommendation( - "storage_container_entity", - Action.OK, - 580_000, - 1_000_000L, - Map.of("autovacuum_vacuum_scale_factor", "0.02"), - Map.of("autovacuum_vacuum_scale_factor", "0.02"), - "ok"))); - - String report = DbTuneReport.render(result); - - assertTrue(report.contains("already match their recommended settings")); - assertFalse(report.contains("Next steps:")); - } - - @Test - void renderAlterStatements_emitsOneSemicolonPerStatement() { - PostgresAutoTuner tuner = new PostgresAutoTuner(); - TableRecommendation a = - new TableRecommendation( - "table_entity", - Action.APPLY, - 500_000, - 1_000_000L, - Map.of(), - Map.of("autovacuum_vacuum_scale_factor", "0.02"), - "ok"); - TableRecommendation b = - new TableRecommendation( - "dashboard_entity", - Action.APPLY, - 300_000, - 1_000_000L, - Map.of(), - Map.of("autovacuum_vacuum_scale_factor", "0.02"), - "ok"); - - String out = DbTuneReport.renderAlterStatements(tuner, List.of(a, b)); - - String[] lines = out.split("\n"); - assertEquals(2, lines.length); - assertTrue(lines[0].endsWith(";")); - assertTrue(lines[1].endsWith(";")); - assertTrue(lines[0].contains("table_entity")); - assertTrue(lines[1].contains("dashboard_entity")); - } - - @Test - void actionableRecommendations_excludesOkAndSkip() { - DbTuneResult result = - new DbTuneResult( - "PostgreSQL", - "17", - List.of(), - List.of( - rec("a", Action.APPLY), - rec("b", Action.OK), - rec("c", Action.SKIP), - rec("d", Action.TIGHTEN), - rec("e", Action.RELAX))); - - List actionable = result.actionableRecommendations(); - - assertEquals(3, actionable.size()); - assertEquals( - List.of("a", "d", "e"), actionable.stream().map(TableRecommendation::tableName).toList()); - } - - private static TableRecommendation rec(final String name, final Action action) { - return new TableRecommendation(name, action, 0, 0, Map.of(), Map.of(), ""); - } -} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/DiagnosticReportTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/DiagnosticReportTest.java deleted file mode 100644 index de2bbbf1964b..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/DiagnosticReportTest.java +++ /dev/null @@ -1,172 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.List; -import java.util.Map; -import org.junit.jupiter.api.Test; - -/** - * Diagnostic-side rendering and grouping tests. Pure logic, no DB. The end-to-end DB query - * exercise lives in {@code DbTuneIT}. - */ -class DiagnosticReportTest { - - @Test - void findingsByCategory_groupsByEnumOrder() { - DbTuneDiagnosis d = - new DbTuneDiagnosis( - List.of( - finding(DiagnosticCategory.SLOW_QUERY, "q1"), - finding(DiagnosticCategory.UNUSED_INDEX, "idx_a"), - finding(DiagnosticCategory.UNUSED_INDEX, "idx_b"), - finding(DiagnosticCategory.HIGH_DEAD_TUPLES, "tag_usage")), - List.of()); - - Map> grouped = d.findingsByCategory(); - - assertEquals(2, grouped.get(DiagnosticCategory.UNUSED_INDEX).size()); - assertEquals(1, grouped.get(DiagnosticCategory.HIGH_DEAD_TUPLES).size()); - assertEquals(1, grouped.get(DiagnosticCategory.SLOW_QUERY).size()); - // EnumMap preserves enum declaration order — UNUSED_INDEX precedes HIGH_DEAD_TUPLES precedes - // SLOW_QUERY. - List orderedKeys = grouped.keySet().stream().toList(); - assertEquals( - List.of( - DiagnosticCategory.UNUSED_INDEX, - DiagnosticCategory.HIGH_DEAD_TUPLES, - DiagnosticCategory.SLOW_QUERY), - orderedKeys); - } - - @Test - void renderDiagnosis_empty_showsCleanResultMessage() { - DbTuneDiagnosis empty = new DbTuneDiagnosis(List.of(), List.of()); - - String out = DbTuneReport.renderDiagnosis(empty); - - assertTrue(out.contains("Diagnostic findings")); - assertTrue(out.contains("every check returned a clean result")); - } - - @Test - void renderDiagnosis_findingsRenderUnderCategorySections() { - DbTuneDiagnosis d = - new DbTuneDiagnosis( - List.of( - new Finding( - DiagnosticCategory.UNUSED_INDEX, - Severity.WARN, - Map.of( - "table", "tag_usage", - "index", "idx_unused_tag", - "size", "120 MB", - "scans", "0"))), - List.of()); - - String out = DbTuneReport.renderDiagnosis(d); - - assertTrue(out.contains("Unused indexes (1 found)")); - assertTrue(out.contains("idx_unused_tag")); - assertTrue(out.contains("120 MB")); - } - - @Test - void renderDiagnosis_notesAppendedWhenPresent() { - DbTuneDiagnosis d = - new DbTuneDiagnosis( - List.of(), List.of("slow queries: pg_stat_statements extension not installed")); - - String out = DbTuneReport.renderDiagnosis(d); - - assertTrue(out.contains("Notes:")); - assertTrue(out.contains("pg_stat_statements extension not installed")); - } - - @Test - void renderDiagnosis_categoriesWithoutFindingsAreSuppressed() { - DbTuneDiagnosis d = - new DbTuneDiagnosis(List.of(finding(DiagnosticCategory.SLOW_QUERY, "SELECT 1")), List.of()); - - String out = DbTuneReport.renderDiagnosis(d); - - assertTrue(out.contains("Top slowest queries")); - assertFalse(out.contains("Unused indexes")); - assertFalse(out.contains("Tables with high dead-tuple ratio")); - } - - @Test - void truncate_collapsesWhitespaceAndAppliesLimit() { - String long_ = - "SELECT *\nFROM table_entity\nWHERE fqnHash LIKE 'foo%' ORDER BY name LIMIT 100"; - - String t = PostgresDiagnostic.truncate(long_); - - assertFalse(t.contains(" ")); - assertFalse(t.contains("\n")); - assertTrue(t.length() <= 101); // 100 + ellipsis - } - - @Test - void truncate_nullReturnsEmpty() { - assertEquals("", PostgresDiagnostic.truncate(null)); - assertEquals("", MysqlDiagnostic.truncate(null)); - } - - @Test - void truncate_underLimitReturnsAsIs() { - assertEquals("SELECT 1", PostgresDiagnostic.truncate("SELECT 1")); - } - - @Test - void truncate_overLimitGetsEllipsis() { - String long_ = "x".repeat(150); - String t = PostgresDiagnostic.truncate(long_); - assertTrue(t.endsWith("…")); - assertEquals(101, t.length()); - } - - @Test - void diagnosticCategory_columnsAreImmutable() { - List cols = DiagnosticCategory.UNUSED_INDEX.columns(); - org.junit.jupiter.api.Assertions.assertThrows( - UnsupportedOperationException.class, () -> cols.add("new_col")); - } - - @Test - void nullSafe_returnsEmptyForNullAndUntouchedForNonNull() { - assertEquals("", PostgresDiagnostic.nullSafe(null)); - assertEquals("", PostgresDiagnostic.nullSafe("")); - assertEquals("2026-05-11 10:00:00", PostgresDiagnostic.nullSafe("2026-05-11 10:00:00")); - } - - @Test - void formatSeqIdxRatio_usesDoubleDivisionAndOneDecimal() { - assertEquals("7.5", PostgresDiagnostic.formatSeqIdxRatio(15, 2)); - assertEquals("10.0", PostgresDiagnostic.formatSeqIdxRatio(100, 10)); - assertEquals("0.5", PostgresDiagnostic.formatSeqIdxRatio(1, 2)); - } - - @Test - void formatSeqIdxRatio_zeroIdxScansRendersInfinity() { - assertEquals("∞", PostgresDiagnostic.formatSeqIdxRatio(50000, 0)); - } - - private static Finding finding(final DiagnosticCategory category, final String objectName) { - return new Finding(category, Severity.INFO, Map.of("table", objectName)); - } -} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/MysqlAutoTunerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/MysqlAutoTunerTest.java deleted file mode 100644 index 22e73875b70a..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/MysqlAutoTunerTest.java +++ /dev/null @@ -1,155 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.Map; -import org.junit.jupiter.api.Test; - -/** - * Heuristic-only tests for the MySQL tuner. Mirrors {@link PostgresAutoTunerTest} but pins the - * MySQL-specific reloption keys (STATS_PERSISTENT / STATS_AUTO_RECALC / STATS_SAMPLE_PAGES) and - * ALTER TABLE syntax (no parens, comma-separated key=value). - */ -class MysqlAutoTunerTest { - - private final MysqlAutoTuner tuner = new MysqlAutoTuner(); - - @Test - void recommend_unknownTable_returnsSkip() { - TableStats stats = stats("not_a_real_table", 1_000_000, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.SKIP, rec.action()); - } - - @Test - void recommend_belowRowThreshold_returnsSkip() { - TableStats stats = stats("storage_container_entity", 100, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.SKIP, rec.action()); - } - - @Test - void recommend_largeEntityWithNoSettings_returnsApply() { - TableStats stats = stats("storage_container_entity", 580_000, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.APPLY, rec.action()); - assertEquals("64", rec.recommendedSettings().get("STATS_SAMPLE_PAGES")); - assertEquals("1", rec.recommendedSettings().get("STATS_PERSISTENT")); - } - - @Test - void recommend_hotTablesGetHigherSampling() { - TableStats stats = stats("tag_usage", 7_400_000, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.APPLY, rec.action()); - assertEquals("100", rec.recommendedSettings().get("STATS_SAMPLE_PAGES")); - } - - @Test - void recommend_alreadyMatching_returnsOk() { - TableStats stats = - stats( - "storage_container_entity", - 580_000, - Map.of( - "STATS_PERSISTENT", "1", - "STATS_AUTO_RECALC", "1", - "STATS_SAMPLE_PAGES", "64")); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.OK, rec.action()); - } - - @Test - void recommend_partialSettings_returnsTighten() { - TableStats stats = - stats("storage_container_entity", 580_000, Map.of("STATS_SAMPLE_PAGES", "20")); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.TIGHTEN, rec.action()); - } - - @Test - void buildAlterStatement_usesMySqlSyntax() { - TableRecommendation rec = - new TableRecommendation( - "storage_container_entity", - Action.APPLY, - 500_000, - 1_000_000_000L, - Map.of(), - Map.of( - "STATS_PERSISTENT", "1", - "STATS_AUTO_RECALC", "1", - "STATS_SAMPLE_PAGES", "64"), - "ok"); - - String sql = tuner.buildAlterStatement(rec); - - assertEquals( - "ALTER TABLE `storage_container_entity` " - + "STATS_AUTO_RECALC=1, STATS_PERSISTENT=1, STATS_SAMPLE_PAGES=64", - sql); - } - - @Test - void parseCreateOptions_emptyAndBlankProduceEmptyMap() { - assertTrue(MysqlAutoTuner.parseCreateOptions(null).isEmpty()); - assertTrue(MysqlAutoTuner.parseCreateOptions("").isEmpty()); - assertTrue(MysqlAutoTuner.parseCreateOptions(" ").isEmpty()); - } - - @Test - void parseCreateOptions_extractsOnlyStatsKeys() { - Map parsed = - MysqlAutoTuner.parseCreateOptions( - "row_format=DYNAMIC stats_persistent=1 stats_sample_pages=64"); - - assertEquals(Map.of("STATS_PERSISTENT", "1", "STATS_SAMPLE_PAGES", "64"), parsed); - } - - @Test - void quoteIdent_usesBacktickAndRejectsUnsafe() { - assertEquals( - "`storage_container_entity`", MysqlAutoTuner.quoteIdent("storage_container_entity")); - assertThrows(IllegalArgumentException.class, () -> MysqlAutoTuner.quoteIdent("`evil`")); - assertThrows(IllegalArgumentException.class, () -> MysqlAutoTuner.quoteIdent("foo;bar")); - } - - @Test - void buildServerCheck_recommendedFormulaIsUntuned() { - ServerParamCheck check = - MysqlAutoTuner.buildServerCheck("innodb_buffer_pool_size", "1073741824", "40-60% of RAM"); - - assertEquals(ServerParamCheck.STATUS_UNTUNED, check.status()); - } - - private static TableStats stats( - final String tableName, final long rowCount, final Map currentSettings) { - return new TableStats(tableName, rowCount, 1_000_000L, 500_000L, currentSettings); - } -} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/PostgresAutoTunerTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/PostgresAutoTunerTest.java deleted file mode 100644 index 3dac30c63ce9..000000000000 --- a/openmetadata-service/src/test/java/org/openmetadata/service/util/dbtune/PostgresAutoTunerTest.java +++ /dev/null @@ -1,251 +0,0 @@ -/* - * Copyright 2026 Collate - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.openmetadata.service.util.dbtune; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.Map; -import org.junit.jupiter.api.Test; - -/** - * Heuristic-only tests — no database. Walks the {@code recommend(stats) -> recommendation} pure - * function across the six action outcomes: SKIP for unknown table, SKIP under threshold, APPLY for - * empty-and-tighten, OK for already-matching, TIGHTEN for partial-match, RELAX for change_event. - * Also pins the SQL-builder format and identifier-quoting safety invariants because both feed - * directly into ALTER TABLE statements. - */ -class PostgresAutoTunerTest { - - private final PostgresAutoTuner tuner = new PostgresAutoTuner(); - - @Test - void recommend_unknownTable_returnsSkip() { - TableStats stats = stats("not_a_real_table", 1_000_000, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.SKIP, rec.action()); - assertTrue(rec.reason().contains("not in the dbtune catalog")); - } - - @Test - void recommend_belowRowThreshold_returnsSkip() { - TableStats stats = stats("storage_container_entity", 50, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.SKIP, rec.action()); - assertTrue(rec.reason().contains("below threshold")); - } - - @Test - void recommend_largeEntityWithNoSettings_returnsApply() { - TableStats stats = stats("storage_container_entity", 580_000, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.APPLY, rec.action()); - assertEquals("0.01", rec.recommendedSettings().get("autovacuum_analyze_scale_factor")); - assertEquals("0.02", rec.recommendedSettings().get("autovacuum_vacuum_scale_factor")); - } - - @Test - void recommend_largeEntityWithLooserSettings_returnsTighten() { - TableStats stats = - stats( - "storage_container_entity", - 580_000, - Map.of( - "autovacuum_analyze_scale_factor", "0.1", - "autovacuum_vacuum_scale_factor", "0.2")); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.TIGHTEN, rec.action()); - } - - @Test - void recommend_alreadyMatching_returnsOk() { - TableStats stats = - stats( - "storage_container_entity", - 580_000, - Map.of( - "autovacuum_analyze_scale_factor", "0.01", - "autovacuum_vacuum_scale_factor", "0.02")); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.OK, rec.action()); - } - - @Test - void recommend_alreadyMatchingNumericallyDifferentTextually_returnsOk() { - TableStats stats = - stats( - "storage_container_entity", - 580_000, - Map.of( - "autovacuum_analyze_scale_factor", "0.010", - "autovacuum_vacuum_scale_factor", "0.0200")); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.OK, rec.action(), "0.010 must equal 0.01 numerically"); - } - - @Test - void recommend_changeEventWithNoSettings_returnsRelax() { - TableStats stats = stats("change_event", 12_000_000, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.RELAX, rec.action()); - assertEquals("0.1", rec.recommendedSettings().get("autovacuum_analyze_scale_factor")); - assertEquals("0.2", rec.recommendedSettings().get("autovacuum_vacuum_scale_factor")); - } - - @Test - void recommend_hotTableHasZeroThreshold() { - TableStats stats = stats("entity_relationship", 1, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.APPLY, rec.action()); - assertEquals("4000", rec.recommendedSettings().get("autovacuum_vacuum_cost_limit")); - } - - @Test - void recommend_tagUsageRecommendsCostDelayZero() { - TableStats stats = stats("tag_usage", 7_400_000, Map.of()); - - TableRecommendation rec = tuner.recommend(stats); - - assertEquals(Action.APPLY, rec.action()); - assertEquals("0", rec.recommendedSettings().get("autovacuum_vacuum_cost_delay")); - } - - @Test - void buildAlterStatement_emitsSortedKeyValuePairs() { - TableRecommendation rec = - new TableRecommendation( - "storage_container_entity", - Action.APPLY, - 500_000, - 1_000_000_000L, - Map.of(), - Map.of( - "autovacuum_analyze_scale_factor", "0.01", - "autovacuum_vacuum_scale_factor", "0.02"), - "ok"); - - String sql = tuner.buildAlterStatement(rec); - - assertEquals( - "ALTER TABLE \"storage_container_entity\" SET (" - + "autovacuum_analyze_scale_factor = 0.01, " - + "autovacuum_vacuum_scale_factor = 0.02)", - sql); - } - - @Test - void parseReloptions_emptyAndNullProduceEmptyMap() { - assertTrue(PostgresAutoTuner.parseReloptions(null).isEmpty()); - assertTrue(PostgresAutoTuner.parseReloptions(new String[0]).isEmpty()); - } - - @Test - void parseReloptions_filtersUnknownKeysAndLowercasesNames() { - Map parsed = - PostgresAutoTuner.parseReloptions( - new String[] {"AUTOVACUUM_VACUUM_SCALE_FACTOR=0.05", "fillfactor=90"}); - - assertEquals(Map.of("autovacuum_vacuum_scale_factor", "0.05"), parsed); - } - - @Test - void quoteIdent_rejectsSqlInjectionAttempts() { - assertThrows( - IllegalArgumentException.class, () -> PostgresAutoTuner.quoteIdent("foo; DROP TABLE bar")); - assertThrows(IllegalArgumentException.class, () -> PostgresAutoTuner.quoteIdent("\"oops\"")); - } - - @Test - void quoteIdent_acceptsValidIdentifiers() { - assertEquals( - "\"storage_container_entity\"", PostgresAutoTuner.quoteIdent("storage_container_entity")); - } - - @Test - void settingsMatch_recommendedSubsetOfCurrent_isMatch() { - Map rec = Map.of("a", "0.01"); - Map current = Map.of("a", "0.01", "b", "999"); - - assertTrue(PostgresAutoTuner.settingsMatch(current, rec)); - } - - @Test - void settingsMatch_missingRecommendedKey_isNotMatch() { - Map rec = Map.of("a", "0.01", "b", "0.02"); - Map current = Map.of("a", "0.01"); - - assertFalse(PostgresAutoTuner.settingsMatch(current, rec)); - } - - @Test - void buildServerCheck_recommendedFormulaIsUntuned() { - ServerParamCheck check = - PostgresAutoTuner.buildServerCheck("shared_buffers", "16384", "40% of RAM"); - - assertEquals(ServerParamCheck.STATUS_UNTUNED, check.status()); - } - - @Test - void buildServerCheck_currentMissingIsUnknown() { - ServerParamCheck check = PostgresAutoTuner.buildServerCheck("missing", null, "200"); - - assertEquals(ServerParamCheck.STATUS_UNKNOWN, check.status()); - } - - @Test - void buildServerCheck_numericMatchIsOk() { - ServerParamCheck check = PostgresAutoTuner.buildServerCheck("random_page_cost", "1.10", "1.1"); - - assertEquals(ServerParamCheck.STATUS_OK, check.status()); - } - - @Test - void buildServerCheck_numericMismatchIsLabelledMismatch() { - ServerParamCheck check = PostgresAutoTuner.buildServerCheck("work_mem", "4096", "131072"); - - assertEquals(ServerParamCheck.STATUS_MISMATCH, check.status()); - } - - @Test - void buildServerCheck_currentHigherThanRecommendedIsAlsoMismatch() { - // random_page_cost recommendation (1.1) is intentionally LOWER than the SSD-naive default - // (4.0). - // Direction-agnostic MISMATCH avoids the misleading "UNDERSIZED" label here. - ServerParamCheck check = PostgresAutoTuner.buildServerCheck("random_page_cost", "4.0", "1.1"); - - assertEquals(ServerParamCheck.STATUS_MISMATCH, check.status()); - } - - private static TableStats stats( - final String tableName, final long rowCount, final Map currentSettings) { - return new TableStats(tableName, rowCount, 1_000_000L, 500_000L, currentSettings); - } -} diff --git a/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdf/customOntology.json b/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdf/customOntology.json new file mode 100644 index 000000000000..be6fa7d6c7cc --- /dev/null +++ b/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdf/customOntology.json @@ -0,0 +1,105 @@ +{ + "$id": "https://open-metadata.org/schema/api/configuration/rdf/customOntology.json", + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "CustomOntology", + "description": "A user-authored extension to the canonical OpenMetadata ontology. Custom classes and properties live in the om-extension: namespace and never collide with the read-only canonical om: namespace.", + "type": "object", + "javaType": "org.openmetadata.schema.api.configuration.rdf.CustomOntology", + "definitions": { + "customClass": { + "type": "object", + "javaType": "org.openmetadata.schema.api.configuration.rdf.CustomOntologyClass", + "description": "A user-defined OWL class.", + "properties": { + "uri": { + "description": "Full URI of the class. Must start with the om-extension: namespace.", + "type": "string", + "pattern": "^https://open-metadata\\.org/ontology-extension/[A-Za-z][A-Za-z0-9_-]*$" + }, + "label": { + "description": "Human-readable label (rdfs:label).", + "type": "string" + }, + "description": { + "description": "Markdown description of the class.", + "type": "string" + }, + "subClassOf": { + "description": "Parent class URIs. May reference canonical om: classes (e.g. om:DataAsset) or other custom classes within this same extension. Must not be empty.", + "type": "array", + "items": { "type": "string" }, + "minItems": 1 + } + }, + "required": ["uri", "subClassOf"], + "additionalProperties": false + }, + "customProperty": { + "type": "object", + "javaType": "org.openmetadata.schema.api.configuration.rdf.CustomOntologyProperty", + "description": "A user-defined OWL ObjectProperty or DatatypeProperty.", + "properties": { + "uri": { + "description": "Full URI of the property. Must start with the om-extension: namespace.", + "type": "string", + "pattern": "^https://open-metadata\\.org/ontology-extension/[A-Za-z][A-Za-z0-9_-]*$" + }, + "label": { + "description": "Human-readable label.", + "type": "string" + }, + "description": { + "description": "Markdown description of the property.", + "type": "string" + }, + "type": { + "description": "OWL property type.", + "type": "string", + "enum": ["ObjectProperty", "DatatypeProperty"] + }, + "domain": { + "description": "URI of the property's rdfs:domain (the class instances this property applies to).", + "type": "string" + }, + "range": { + "description": "URI of the property's rdfs:range. For DatatypeProperty, an XSD datatype URI; for ObjectProperty, a class URI.", + "type": "string" + }, + "subPropertyOf": { + "description": "Optional parent properties.", + "type": "array", + "items": { "type": "string" }, + "default": [] + } + }, + "required": ["uri", "type", "domain", "range"], + "additionalProperties": false + } + }, + "properties": { + "name": { + "description": "Stable identifier for the extension. Lowercase letters, digits, hyphen.", + "type": "string", + "pattern": "^[a-z][a-z0-9-]{1,62}[a-z0-9]$" + }, + "displayName": { "type": "string" }, + "description": { + "description": "Markdown description of the extension. Should explain why these classes/properties are needed.", + "type": "string" + }, + "classes": { + "description": "Custom OWL classes defined by this extension.", + "type": "array", + "items": { "$ref": "#/definitions/customClass" }, + "default": [] + }, + "properties": { + "description": "Custom OWL properties defined by this extension.", + "type": "array", + "items": { "$ref": "#/definitions/customProperty" }, + "default": [] + } + }, + "required": ["name"], + "additionalProperties": false +} diff --git a/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdf/inferenceRule.json b/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdf/inferenceRule.json new file mode 100644 index 000000000000..405daafc5064 --- /dev/null +++ b/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdf/inferenceRule.json @@ -0,0 +1,54 @@ +{ + "$id": "https://open-metadata.org/schema/api/configuration/rdf/inferenceRule.json", + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "InferenceRule", + "description": "A SPARQL CONSTRUCT rule that materializes derived triples in the OpenMetadata knowledge graph (e.g. transitive lineage, PII propagation, tag inheritance).", + "type": "object", + "javaType": "org.openmetadata.schema.api.configuration.rdf.InferenceRule", + "properties": { + "name": { + "description": "Stable identifier for the rule (used as primary key). Lowercase letters, digits, hyphen.", + "type": "string", + "pattern": "^[a-z][a-z0-9-]{1,62}[a-z0-9]$" + }, + "displayName": { + "description": "Human-readable name.", + "type": "string" + }, + "description": { + "description": "What the rule does and why it is enabled. Markdown.", + "type": "string" + }, + "ruleType": { + "description": "Body language. CONSTRUCT is a SPARQL CONSTRUCT query that produces new triples. RDFS is a placeholder for future Jena-RDFS rule format.", + "type": "string", + "enum": ["CONSTRUCT", "RDFS"], + "default": "CONSTRUCT" + }, + "ruleBody": { + "description": "The rule body. For ruleType=CONSTRUCT, a SPARQL CONSTRUCT query that emits the inferred triples.", + "type": "string", + "minLength": 16 + }, + "enabled": { + "description": "Whether the rule is currently active. Disabled rules are loaded but not applied.", + "type": "boolean", + "default": true + }, + "priority": { + "description": "Execution order hint. Lower numbers run first. Rules at the same priority run in name order.", + "type": "integer", + "default": 100, + "minimum": 0, + "maximum": 10000 + }, + "tags": { + "description": "Free-form labels (e.g. 'lineage', 'security', 'governance') for filtering in admin UI.", + "type": "array", + "items": { "type": "string" }, + "default": [] + } + }, + "required": ["name", "ruleBody"], + "additionalProperties": false +} diff --git a/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdfConfiguration.json b/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdfConfiguration.json index c868823c0465..b2102fa5560c 100644 --- a/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdfConfiguration.json +++ b/openmetadata-spec/src/main/resources/json/schema/api/configuration/rdfConfiguration.json @@ -63,6 +63,29 @@ "description": "Cache inferred triples for better query performance (requires more storage)", "type": "boolean", "default": false + }, + "federation": { + "description": "Controls federated SPARQL access (SERVICE clauses) to external endpoints. Federation is disabled by default; SERVICE clauses are rejected unless the target URI is in the allowlist.", + "type": "object", + "javaType": "org.openmetadata.schema.api.configuration.rdf.SparqlFederationConfig", + "properties": { + "enabled": { + "description": "Master switch for federated SPARQL. When false, every SERVICE clause is rejected regardless of allowlist contents.", + "type": "boolean", + "default": false + }, + "allowedEndpoints": { + "description": "External SPARQL endpoint URIs that may appear in SERVICE clauses. Compared verbatim against the URI in the SERVICE clause; trailing slashes matter.", + "type": "array", + "items": { + "type": "string", + "format": "uri" + }, + "default": [] + } + }, + "additionalProperties": false, + "default": null } }, "required": ["enabled", "storageType"], diff --git a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/chartFunctions.json b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/chartFunctions.json deleted file mode 100644 index 0a3d712a8d4c..000000000000 --- a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/chartFunctions.json +++ /dev/null @@ -1,41 +0,0 @@ -{ - "$id": "https://open-metadata.org/schema/dataInsight/custom/chartFunctions.json", - "$schema": "http://json-schema.org/draft-07/schema#", - "title": "ChartFunctions", - "description": "Shared aggregation function and KPI types referenced by data insight chart configurations. Kept in its own schema so lineChart/summaryCard/formulaHolder/dataInsightCustomChart can $ref these without forming a circular import in the generated code.", - "type": "object", - "definitions": { - "function": { - "javaType": "org.openmetadata.schema.dataInsight.custom.Function", - "description": "aggregation function for chart", - "type": "string", - "enum": [ - "count", - "sum", - "avg", - "min", - "max", - "unique" - ] - }, - "kpiDetails": { - "type": "object", - "javaType": "org.openmetadata.schema.dataInsight.custom.KPIDetails", - "description": "KPI details for the data insight chart.", - "properties": { - "startDate": { - "description": "Start Date of KPI", - "type": "string" - }, - "endDate": { - "description": "End Date of KPI", - "type": "string" - }, - "target": { - "description": "Target value of KPI", - "type": "number" - } - } - } - } -} diff --git a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChart.json b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChart.json index 9fa968107ca7..9af4a0c21a5b 100644 --- a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChart.json +++ b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChart.json @@ -8,6 +8,40 @@ "javaInterfaces": [ "org.openmetadata.schema.EntityInterface" ], + "definitions": { + "function": { + "javaType": "org.openmetadata.schema.dataInsight.custom.Function", + "description": "aggregation function for chart", + "type": "string", + "enum": [ + "count", + "sum", + "avg", + "min", + "max", + "unique" + ] + }, + "kpiDetails": { + "type": "object", + "javaType": "org.openmetadata.schema.dataInsight.custom.KPIDetails", + "description": "KPI details for the data insight chart.", + "properties": { + "startDate": { + "description": "Start Date of KPI", + "type": "string" + }, + "endDate": { + "description": "End Date of KPI", + "type": "string" + }, + "target": { + "description": "Target value of KPI", + "type": "number" + } + } + } + }, "properties": { "id": { "description": "Unique identifier of this table instance.", diff --git a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChartResultList.json b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChartResultList.json index 5f095652609e..1e71481c5f67 100644 --- a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChartResultList.json +++ b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/dataInsightCustomChartResultList.json @@ -14,7 +14,7 @@ } }, "kpiDetails": { - "$ref": "chartFunctions.json#/definitions/kpiDetails" + "$ref": "dataInsightCustomChart.json#/definitions/kpiDetails" } }, "additionalProperties": false diff --git a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/formulaHolder.json b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/formulaHolder.json index b47b10f38450..30d68325d3e6 100644 --- a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/formulaHolder.json +++ b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/formulaHolder.json @@ -11,7 +11,7 @@ "type": "string" }, "function": { - "$ref": "chartFunctions.json#/definitions/function" + "$ref": "dataInsightCustomChart.json#/definitions/function" }, "field": { "description": "Group of Result", diff --git a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/lineChart.json b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/lineChart.json index 40f6bd664204..844203c24121 100644 --- a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/lineChart.json +++ b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/lineChart.json @@ -16,7 +16,7 @@ "type": "string" }, "function": { - "$ref": "chartFunctions.json#/definitions/function" + "$ref": "dataInsightCustomChart.json#/definitions/function" }, "field": { "description": "Filter field for the data insight chart.", @@ -80,7 +80,7 @@ "type": "string" }, "kpiDetails": { - "$ref": "chartFunctions.json#/definitions/kpiDetails" + "$ref": "dataInsightCustomChart.json#/definitions/kpiDetails" }, "xAxisField": { "description": "X-axis field for the data insight chart.", diff --git a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/summaryCard.json b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/summaryCard.json index 3e02e4885faf..8491e3b843b4 100644 --- a/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/summaryCard.json +++ b/openmetadata-spec/src/main/resources/json/schema/dataInsight/custom/summaryCard.json @@ -16,7 +16,7 @@ "type": "string" }, "function": { - "$ref": "chartFunctions.json#/definitions/function" + "$ref": "dataInsightCustomChart.json#/definitions/function" }, "field": { "description": "Filter field for the data insight chart.", diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json b/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json index c37fbd5e5e02..a509d925042e 100644 --- a/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json +++ b/openmetadata-spec/src/main/resources/json/schema/entity/applications/configuration/internal/searchIndexingAppConfig.json @@ -28,6 +28,11 @@ "default": ["all"], "uniqueItems": true }, + "recreateIndex": { + "description": "This schema publisher run modes.", + "type": "boolean", + "default": true + }, "batchSize": { "description": "Maximum number of events sent in a batch (Default 100).", "type": "integer", @@ -82,7 +87,7 @@ "default": 100 }, "searchIndexMappingLanguage": { - "description": "Search index mapping language.", + "description": "Recreate Indexes with updated Language", "$ref": "../../../../configuration/elasticSearchConfiguration.json#/definitions/searchIndexMappingLanguage" }, "autoTune": { @@ -91,6 +96,12 @@ "type": "boolean", "default": false }, + "useDistributedIndexing": { + "title": "Use Distributed Indexing", + "description": "Enable distributed indexing to scale reindexing across multiple servers with fault tolerance and parallel processing", + "type": "boolean", + "default": true + }, "partitionSize": { "title": "Partition Size", "description": "Number of entities per partition for distributed indexing. Smaller values create more partitions for better distribution across servers. Range: 1000-50000.", diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/applications/mcp/mcpToolCallUsage.json b/openmetadata-spec/src/main/resources/json/schema/entity/applications/mcp/mcpToolCallUsage.json deleted file mode 100644 index ce02a7b7fcb0..000000000000 --- a/openmetadata-spec/src/main/resources/json/schema/entity/applications/mcp/mcpToolCallUsage.json +++ /dev/null @@ -1,40 +0,0 @@ -{ - "$id": "https://open-metadata.org/schema/entity/applications/mcp/mcpToolCallUsage.json", - "$schema": "http://json-schema.org/draft-07/schema#", - "title": "McpToolCallUsage", - "description": "Single MCP tool-call usage record. One row written per tool invocation to the apps_extension_time_series table with extension='limits' (reusing the existing per-app usage extension; rows are isolated by appName='McpApplication'). Used to surface MCP traffic as a product growth metric. Not billed, no enforcement.", - "type": "object", - "javaType": "org.openmetadata.schema.entity.app.mcp.McpToolCallUsage", - "properties": { - "appId": { - "description": "Unique identifier of the McpApplication.", - "$ref": "../../../type/basic.json#/definitions/uuid" - }, - "appName": { - "description": "Name of the application (McpApplication).", - "$ref": "../../../type/basic.json#/definitions/entityName" - }, - "timestamp": { - "description": "Time the tool call completed (epoch millis, UTC).", - "$ref": "../../../type/basic.json#/definitions/timestamp" - }, - "extension": { - "$ref": "../../applications/appExtension.json#/definitions/extensionType", - "default": "limits" - }, - "toolName": { - "description": "Name of the MCP tool that was invoked (e.g. search_metadata, create_glossary, nlq_search).", - "type": "string" - }, - "userName": { - "description": "Principal name from the MCP request's security context.", - "type": "string" - }, - "success": { - "description": "True if the tool call returned without an error result.", - "type": "boolean" - } - }, - "additionalProperties": false, - "required": ["extension"] -} diff --git a/openmetadata-spec/src/main/resources/rdf/contexts/ai.jsonld b/openmetadata-spec/src/main/resources/rdf/contexts/ai.jsonld new file mode 100644 index 000000000000..67f6bd883dda --- /dev/null +++ b/openmetadata-spec/src/main/resources/rdf/contexts/ai.jsonld @@ -0,0 +1,97 @@ +{ + "@context": [ + "./base.jsonld", + { + "LLMModel": { + "@id": "om:LLMModel", + "@type": "om:Entity" + }, + "AIApplication": { + "@id": "om:AIApplication", + "@type": "om:Entity" + }, + "McpServer": { + "@id": "om:McpServer", + "@type": "om:Entity" + }, + "AgentExecution": { + "@id": "om:AgentExecution", + "@type": ["om:Entity", "prov:Activity"] + }, + "McpExecution": { + "@id": "om:McpExecution", + "@type": ["om:Entity", "prov:Activity"] + }, + "PromptTemplate": { + "@id": "om:PromptTemplate", + "@type": "om:Entity" + }, + "modelType": { + "@id": "om:modelType", + "@type": "xsd:string" + }, + "modelProvider": { + "@id": "om:modelProvider", + "@type": "xsd:string" + }, + "modelVersion": { + "@id": "om:modelVersion", + "@type": "xsd:string" + }, + "applicationType": { + "@id": "om:applicationType", + "@type": "xsd:string" + }, + "developmentStage": { + "@id": "om:developmentStage", + "@type": "xsd:string" + }, + "trainingDatasets": { + "@id": "om:hasTrainingDataset", + "@type": "@id", + "@container": "@set" + }, + "validationDatasets": { + "@id": "om:hasValidationDataset", + "@type": "@id", + "@container": "@set" + }, + "models": { + "@id": "om:usesModel", + "@type": "@id", + "@container": "@set" + }, + "mcpServers": { + "@id": "om:usesMcpServer", + "@type": "@id", + "@container": "@set" + }, + "tools": { + "@id": "om:usesTool", + "@type": "@id", + "@container": "@set" + }, + "promptTemplates": { + "@id": "om:hasPromptTemplate", + "@type": "@id", + "@container": "@set" + }, + "application": { + "@id": "om:executedBy", + "@type": "@id" + }, + "executionStatus": { + "@id": "om:executionStatus", + "@type": "xsd:string" + }, + "startTime": { + "@id": "prov:startedAtTime", + "@type": "xsd:dateTime" + }, + "endTime": { + "@id": "prov:endedAtTime", + "@type": "xsd:dateTime" + } + } + ] +} diff --git a/openmetadata-spec/src/main/resources/rdf/contexts/automation.jsonld b/openmetadata-spec/src/main/resources/rdf/contexts/automation.jsonld new file mode 100644 index 000000000000..46691e77c992 --- /dev/null +++ b/openmetadata-spec/src/main/resources/rdf/contexts/automation.jsonld @@ -0,0 +1,47 @@ +{ + "@context": [ + "./base.jsonld", + { + "Workflow": { + "@id": "om:Workflow", + "@type": "om:Entity" + }, + "Automation": { + "@id": "om:Automation", + "@type": "om:Entity" + }, + "WorkflowDefinition": { + "@id": "om:WorkflowDefinition", + "@type": "om:Entity" + }, + "WorkflowInstance": { + "@id": "om:WorkflowInstance", + "@type": ["om:Entity", "prov:Activity"] + }, + "workflowType": { + "@id": "om:workflowType", + "@type": "xsd:string" + }, + "status": { + "@id": "om:hasStatus", + "@type": "xsd:string" + }, + "request": { + "@id": "om:automationRequest", + "@type": "@json" + }, + "response": { + "@id": "om:automationResponse", + "@type": "@json" + }, + "triggerType": { + "@id": "om:triggerType", + "@type": "xsd:string" + }, + "scheduleType": { + "@id": "om:scheduleType", + "@type": "xsd:string" + } + } + ] +} diff --git a/openmetadata-spec/src/main/resources/rdf/contexts/governance.jsonld b/openmetadata-spec/src/main/resources/rdf/contexts/governance.jsonld index a87d08e14353..08da94224fce 100644 --- a/openmetadata-spec/src/main/resources/rdf/contexts/governance.jsonld +++ b/openmetadata-spec/src/main/resources/rdf/contexts/governance.jsonld @@ -35,9 +35,18 @@ "@type": "om:Entity" }, "glossary": { - "@id": "om:belongsToGlossary", + "@id": "skos:inScheme", "@type": "@id" }, + "parent": { + "@id": "skos:broader", + "@type": "@id" + }, + "children": { + "@id": "skos:narrower", + "@type": "@id", + "@container": "@set" + }, "synonyms": { "@id": "skos:altLabel", "@type": "xsd:string", @@ -62,7 +71,7 @@ "@type": "xsd:string" }, "classification": { - "@id": "om:belongsToClassification", + "@id": "skos:inScheme", "@type": "@id" }, "usageCount": { @@ -86,6 +95,32 @@ "@id": "om:contractForEntity", "@type": "xsd:string" }, + "contractStatus": { + "@id": "om:contractStatus", + "@type": "xsd:string" + }, + "appliesTo": { + "@id": "om:appliesToEntity", + "@type": "@id" + }, + "schema": { + "@id": "om:hasContractTerm", + "@type": "@json" + }, + "qualityExpectations": { + "@id": "om:hasContractTerm", + "@type": "@json", + "@container": "@set" + }, + "slaExpectations": { + "@id": "om:hasContractTerm", + "@type": "@json" + }, + "users": { + "@id": "om:appliesToUser", + "@type": "@id", + "@container": "@set" + }, "domainType": { "@id": "om:domainType", "@type": "xsd:string" @@ -130,11 +165,6 @@ "@type": "@id", "@container": "@set" }, - "childTerms": { - "@id": "skos:narrower", - "@type": "@id", - "@container": "@set" - }, "termType": { "@id": "om:termType", "@type": "xsd:string" diff --git a/openmetadata-spec/src/main/resources/rdf/contexts/lineage.jsonld b/openmetadata-spec/src/main/resources/rdf/contexts/lineage.jsonld index 540479340277..096285789d98 100644 --- a/openmetadata-spec/src/main/resources/rdf/contexts/lineage.jsonld +++ b/openmetadata-spec/src/main/resources/rdf/contexts/lineage.jsonld @@ -71,12 +71,12 @@ "@container": "@set" }, "fromColumns": { - "@id": "om:fromColumns", + "@id": "om:fromColumnFqn", "@type": "xsd:string", "@container": "@list" }, "toColumn": { - "@id": "om:toColumn", + "@id": "om:toColumnFqn", "@type": "xsd:string" }, "function": { diff --git a/openmetadata-spec/src/main/resources/rdf/ontology/CHANGELOG.md b/openmetadata-spec/src/main/resources/rdf/ontology/CHANGELOG.md new file mode 100644 index 000000000000..8e9964d57a39 --- /dev/null +++ b/openmetadata-spec/src/main/resources/rdf/ontology/CHANGELOG.md @@ -0,0 +1,62 @@ +# OpenMetadata Ontology Changelog + +The canonical ontology lives in `openmetadata.ttl`. The PROV-aligned extension lives in +`openmetadata-prov.ttl`. SHACL shapes live in `../shapes/openmetadata-shapes.ttl`. JSON-LD contexts +live in `../contexts/`. + +The version recorded here is the value of `owl:versionInfo` on the `om:` ontology resource. + +## 1.1.0 — 2026-04-28 + +Knowledge-graph fidelity pass. All changes are additive or domain corrections; existing consumers +that referenced the corrected domains were not actually relying on them, since the prior +declaration did not match what the mapper emitted. + +### Added — Column resources and column lineage + +- `om:Column` resources are now first-class named resources at FQN-derived URIs + (`baseUri + "entity/column/" + URLEncoded(FQN)`). Previously columns were blank nodes, + unreachable from SPARQL. +- `om:fromColumn` and `om:toColumn` (column lineage) are now URI references to `om:Column` + resources, not FQN string literals. The original FQN strings are retained as + `om:fromColumnFqn` / `om:toColumnFqn` for back-compatibility with consumers that match + by string. +- `om:LineageDetails` class declared (was used by the mapper but undeclared). +- `om:hasColumnLineage`, `om:transformFunction` declared. +- `om:hasChildColumn` (subproperty of `om:hasColumn`) for nested struct/map/union columns. +- Domain of `om:fromColumn` / `om:toColumn` corrected from `om:Column` to `om:ColumnLineage`. + +### Added — Table constraints + +- `om:TableConstraint` class. +- `om:hasConstraint` (`om:Table` → `om:TableConstraint`). +- `om:constraintType`, `om:relationshipType` (datatype properties). +- `om:hasConstrainedColumn`, `om:hasReferredColumn` (object properties on the constraint). +- `om:references` (`om:Column` → `om:Column`) — direct FK edges between source and referred + columns, paired positionally from `TableConstraint.columns` and `referredColumns`. +- `om:isUnique` datatype property on columns. +- Per-column `constraint` enum (`PRIMARY_KEY`, `UNIQUE`, `NOT_NULL`, `NULL`) now maps to the + corresponding `om:isPrimaryKey` / `om:isUnique` / `om:isNullable` triples. + +### Changed — SKOS hierarchy + +JSON-LD context `governance.jsonld`: + +- `glossary` (on a glossary term) now maps to `skos:inScheme` (was `om:belongsToGlossary`). +- `classification` (on a tag) now maps to `skos:inScheme` (was `om:belongsToClassification`). +- `parent` (on a glossary term or tag) now maps to `skos:broader` (previously unmapped). +- `children` now maps to `skos:narrower`. The prior `childTerms` alias has been removed; it + referenced a JSON field that does not exist on `GlossaryTerm`, so the mapping never fired. + +The OpenMetadata-specific predicates `om:belongsToGlossary` and `om:belongsToClassification` +were not used outside this single context file; no SPARQL queries reference them. + +### Changed — JSON-LD lineage context + +- `fromColumns` and `toColumn` in `lineage.jsonld` now map to `om:fromColumnFqn` / + `om:toColumnFqn` (datatype properties), not to `om:fromColumns` / `om:toColumn` (which + collided with the new object-typed predicates). + +## 1.0.0 — 2025-08-24 + +Initial ontology release. diff --git a/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata-prov.ttl b/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata-prov.ttl index dc7d781c1218..3ab0f799452a 100644 --- a/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata-prov.ttl +++ b/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata-prov.ttl @@ -1,6 +1,7 @@ @prefix om: . @prefix prov: . @prefix xsd: . +@prefix rdf: . @prefix rdfs: . @prefix dcterms: . diff --git a/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata.ttl b/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata.ttl index 074798175812..3a89a34d9017 100644 --- a/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata.ttl +++ b/openmetadata-spec/src/main/resources/rdf/ontology/openmetadata.ttl @@ -24,13 +24,15 @@ # OpenMetadata Complete Ontology om: a owl:Ontology ; - owl:versionInfo "1.0.0" ; + owl:versionInfo "1.1.0" ; + owl:priorVersion ; dct:title "OpenMetadata Ontology" ; dct:description "Complete ontology for OpenMetadata covering all entities and relationships" ; dct:creator "OpenMetadata Team" ; dct:created "2025-08-24"^^xsd:date ; + dct:modified "2026-04-28"^^xsd:date ; dct:license ; - owl:imports dcat: , prov: , skos: . + owl:imports dcat: , prov: , skos: , dqv: . ################################################################# # Base Classes @@ -315,8 +317,27 @@ om:usedToCalculate a owl:ObjectProperty ; om:DataContract a owl:Class ; rdfs:label "Data Contract" ; + rdfs:comment "A formal agreement defining the schema, semantics, quality expectations, and SLAs for a data asset." ; rdfs:subClassOf om:Entity . +om:contractStatus a owl:DatatypeProperty ; + rdfs:label "contract status" ; + rdfs:comment "Lifecycle stage of the contract (Draft, Active, Deprecated)." ; + rdfs:domain om:DataContract ; + rdfs:range xsd:string . + +om:appliesToEntity a owl:ObjectProperty ; + rdfs:label "applies to entity" ; + rdfs:comment "The data asset this contract governs." ; + rdfs:domain om:DataContract ; + rdfs:range om:Entity . + +om:hasContractTerm a owl:ObjectProperty ; + rdfs:label "has contract term" ; + rdfs:comment "Schema, quality, semantic, or SLA term embedded in this contract." ; + rdfs:domain om:DataContract ; + rdfs:range om:Entity . + om:DataProduct a owl:Class ; rdfs:label "Data Product" ; rdfs:subClassOf om:Entity . @@ -325,6 +346,142 @@ om:Domain a owl:Class ; rdfs:label "Domain" ; rdfs:subClassOf om:Entity . +om:Persona a owl:Class ; + rdfs:label "Persona" ; + rdfs:comment "A role or persona that aggregates users with similar needs and customized UI experiences." ; + rdfs:subClassOf om:Entity . + +################################################################# +# AI & LLM Classes +################################################################# + +om:LLMModel a owl:Class ; + rdfs:label "LLM Model" ; + rdfs:comment "A registered Large Language Model deployment, fine-tune, or base model." ; + rdfs:subClassOf om:Entity . + +om:AIApplication a owl:Class ; + rdfs:label "AI Application" ; + rdfs:comment "An AI system: chatbot, agent, copilot, RAG application, etc." ; + rdfs:subClassOf om:Entity . + +om:McpServer a owl:Class ; + rdfs:label "MCP Server" ; + rdfs:comment "Model Context Protocol server exposing tools and data sources to AI agents." ; + rdfs:subClassOf om:Entity . + +om:AgentExecution a owl:Class ; + rdfs:label "Agent Execution" ; + rdfs:comment "A single run of an AI agent or application." ; + rdfs:subClassOf om:Entity, prov:Activity . + +om:McpExecution a owl:Class ; + rdfs:label "MCP Execution" ; + rdfs:comment "A single tool invocation against an MCP server." ; + rdfs:subClassOf om:Entity, prov:Activity . + +om:PromptTemplate a owl:Class ; + rdfs:label "Prompt Template" ; + rdfs:comment "Reusable prompt template for AI applications." ; + rdfs:subClassOf om:Entity . + +om:hasTrainingDataset a owl:ObjectProperty ; + rdfs:label "has training dataset" ; + rdfs:comment "Dataset used to train or fine-tune an LLM. Critical for AI data lineage." ; + rdfs:domain om:LLMModel ; + rdfs:range om:DataAsset . + +om:hasValidationDataset a owl:ObjectProperty ; + rdfs:label "has validation dataset" ; + rdfs:domain om:LLMModel ; + rdfs:range om:DataAsset . + +om:usesModel a owl:ObjectProperty ; + rdfs:label "uses model" ; + rdfs:comment "An AI application uses one or more LLM models." ; + rdfs:domain om:AIApplication ; + rdfs:range om:LLMModel . + +om:usesMcpServer a owl:ObjectProperty ; + rdfs:label "uses MCP server" ; + rdfs:domain om:AIApplication ; + rdfs:range om:McpServer . + +################################################################# +# Data Quality Metrics (DQV-aligned) +################################################################# + +om:DataQualityMetric a owl:Class ; + rdfs:label "Data Quality Metric" ; + rdfs:comment "Base class for OpenMetadata data-quality metrics. Each instance is a specific named metric (e.g. om:RowCountMetric) emitted as the dqv:isMeasurementOf target on dqv:QualityMeasurement instances." ; + rdfs:subClassOf dqv:Metric . + +# Table-level metrics +om:RowCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "row count" . +om:ColumnCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "column count" . +om:SizeInBytesMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "size in bytes" . + +# Column-level cardinality metrics +om:ValuesCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "values count" . +om:ValidCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "valid count" . +om:NullCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "null count" . +om:NullProportionMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "null proportion" . +om:MissingCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "missing count" . +om:MissingPercentageMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "missing percentage" . +om:UniqueCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "unique count" . +om:UniqueProportionMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "unique proportion" . +om:DistinctCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "distinct count" . +om:DistinctProportionMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "distinct proportion" . +om:DuplicateCountMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "duplicate count" . + +# Column-level statistical metrics +om:MeanMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "mean" . +om:MinMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "min" . +om:MaxMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "max" . +om:MinLengthMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "min length" . +om:MaxLengthMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "max length" . +om:SumMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "sum" . +om:StddevMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "stddev" . +om:VarianceMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "variance" . +om:MedianMetric a om:DataQualityMetric, dqv:Metric ; + rdfs:label "median" . + +################################################################# +# Automation Classes +################################################################# + +om:Workflow a owl:Class ; + rdfs:label "Workflow" ; + rdfs:comment "A configured automation workflow (e.g. ingestion pipeline definition)." ; + rdfs:subClassOf om:Entity . + +om:Automation a owl:Class ; + rdfs:label "Automation" ; + rdfs:comment "An automation request (e.g. test connection, run query) that produces a result." ; + rdfs:subClassOf om:Entity . + ################################################################# # Teams & Users Classes ################################################################# @@ -604,6 +761,11 @@ om:Column a owl:Class ; rdfs:comment "Column in a table" ; rdfs:subClassOf om:Entity . +om:LineageDetails a owl:Class ; + rdfs:label "Lineage Details" ; + rdfs:comment "Metadata describing a single lineage edge between two data assets, including SQL, column-level mapping, and the pipeline that produced it." ; + rdfs:subClassOf prov:Derivation . + om:hasColumn a owl:ObjectProperty ; rdfs:label "has column" ; rdfs:domain om:Table ; @@ -1100,6 +1262,94 @@ om:modified a owl:DatatypeProperty ; rdfs:domain om:Entity ; rdfs:range xsd:dateTime . +################################################################# +# Usage Properties (sourced from UsageDetails on Table / Dashboard / etc.) +################################################################# + +om:usageDailyCount a owl:DatatypeProperty ; + rdfs:label "usage daily count" ; + rdfs:comment "Number of times this asset was queried on the recorded date." ; + rdfs:domain om:Entity ; + rdfs:range xsd:long . + +om:usageDailyPercentile a owl:DatatypeProperty ; + rdfs:label "usage daily percentile" ; + rdfs:comment "Daily percentile rank of this asset's query usage among assets of the same type. Range 0-100. Used by /v1/rdf/insights/important to rank entities." ; + rdfs:domain om:Entity ; + rdfs:range xsd:double . + +om:usageWeeklyCount a owl:DatatypeProperty ; + rdfs:label "usage weekly count" ; + rdfs:comment "Rolling 7-day query count." ; + rdfs:domain om:Entity ; + rdfs:range xsd:long . + +om:usageWeeklyPercentile a owl:DatatypeProperty ; + rdfs:label "usage weekly percentile" ; + rdfs:domain om:Entity ; + rdfs:range xsd:double . + +om:usageMonthlyCount a owl:DatatypeProperty ; + rdfs:label "usage monthly count" ; + rdfs:comment "Rolling 30-day query count." ; + rdfs:domain om:Entity ; + rdfs:range xsd:long . + +om:usageMonthlyPercentile a owl:DatatypeProperty ; + rdfs:label "usage monthly percentile" ; + rdfs:domain om:Entity ; + rdfs:range xsd:double . + +om:usageDate a owl:DatatypeProperty ; + rdfs:label "usage date" ; + rdfs:comment "Date the usage stats were recorded for." ; + rdfs:domain om:Entity ; + rdfs:range xsd:date . + +################################################################# +# Insights (Phase 3 — graph-derived importance scores) +################################################################# + +om:centralityScore a owl:DatatypeProperty ; + rdfs:label "centrality score" ; + rdfs:comment "Topology-derived importance score (e.g. weighted PageRank) computed for entities lacking query usage. Range 0-1. Persisted into a separate named graph by the RdfGraphAlgorithmsApp." ; + rdfs:domain om:Entity ; + rdfs:range xsd:double . + +om:centralityRank a owl:DatatypeProperty ; + rdfs:label "centrality rank" ; + rdfs:comment "Rank within the entity type, lowest is most central." ; + rdfs:domain om:Entity ; + rdfs:range xsd:integer . + +om:Community a owl:Class ; + rdfs:label "Community" ; + rdfs:comment "A cluster of entities discovered by community-detection over the lineage or tag-co-occurrence graph (e.g. Louvain modularity optimization). Members share dense connectivity within the cluster relative to the rest of the graph." . + +om:hasMember a owl:ObjectProperty ; + rdfs:label "has member" ; + rdfs:comment "Asserts that an entity belongs to a community." ; + rdfs:domain om:Community ; + rdfs:range om:Entity . + +om:modularity a owl:DatatypeProperty ; + rdfs:label "modularity" ; + rdfs:comment "Modularity score (Newman, 2006) of the partition that produced this community. Higher values indicate more clearly defined communities; the same value applies to every community persisted from the same run. Range [-1, 1]." ; + rdfs:domain om:Community ; + rdfs:range xsd:double . + +om:communityType a owl:DatatypeProperty ; + rdfs:label "community type" ; + rdfs:comment "Source graph the community was discovered on, e.g. 'lineage' or 'tagCoOccurrence'." ; + rdfs:domain om:Community ; + rdfs:range xsd:string . + +om:communitySize a owl:DatatypeProperty ; + rdfs:label "community size" ; + rdfs:comment "Number of members in this community." ; + rdfs:domain om:Community ; + rdfs:range xsd:integer . + ################################################################# # CSVW Tabular Data Metadata (for Tables) ################################################################# @@ -1142,6 +1392,57 @@ om:isNullable a owl:DatatypeProperty ; rdfs:domain om:Column ; rdfs:range xsd:boolean . +om:isUnique a owl:DatatypeProperty ; + rdfs:label "is unique" ; + rdfs:comment "Whether this column has a uniqueness guarantee, either from a column-level UNIQUE constraint or membership in a single-column UNIQUE / PRIMARY_KEY table constraint." ; + rdfs:domain om:Column ; + rdfs:range xsd:boolean . + +om:references a owl:ObjectProperty ; + rdfs:label "references" ; + rdfs:comment "Foreign-key reference: this column references the value space of another column. The triple is positional — source columns[i] references referredColumns[i] from the table constraint." ; + rdfs:domain om:Column ; + rdfs:range om:Column . + +################################################################# +# Table-level Constraint Properties +################################################################# + +om:TableConstraint a owl:Class ; + rdfs:label "Table Constraint" ; + rdfs:comment "A table-level constraint such as PRIMARY_KEY, UNIQUE, or FOREIGN_KEY that may span multiple columns." ; + rdfs:subClassOf om:Entity . + +om:hasConstraint a owl:ObjectProperty ; + rdfs:label "has constraint" ; + rdfs:comment "Links a table to a constraint declared on it." ; + rdfs:domain om:Table ; + rdfs:range om:TableConstraint . + +om:constraintType a owl:DatatypeProperty ; + rdfs:label "constraint type" ; + rdfs:comment "PRIMARY_KEY, UNIQUE, FOREIGN_KEY, SORT_KEY, DIST_KEY, or CLUSTER_KEY." ; + rdfs:domain om:TableConstraint ; + rdfs:range xsd:string . + +om:hasConstrainedColumn a owl:ObjectProperty ; + rdfs:label "has constrained column" ; + rdfs:comment "A column participating in this constraint (the source side for a FOREIGN_KEY)." ; + rdfs:domain om:TableConstraint ; + rdfs:range om:Column . + +om:hasReferredColumn a owl:ObjectProperty ; + rdfs:label "has referred column" ; + rdfs:comment "A column that this constraint references (only set for FOREIGN_KEY)." ; + rdfs:domain om:TableConstraint ; + rdfs:range om:Column . + +om:relationshipType a owl:DatatypeProperty ; + rdfs:label "relationship type" ; + rdfs:comment "Cardinality of a FOREIGN_KEY constraint: ONE_TO_ONE, ONE_TO_MANY, MANY_TO_ONE, or MANY_TO_MANY." ; + rdfs:domain om:TableConstraint ; + rdfs:range xsd:string . + om:sampleData a owl:ObjectProperty ; rdfs:label "sample data" ; rdfs:comment "Sample data from the table" ; @@ -1173,17 +1474,48 @@ om:transformationFunction a owl:DatatypeProperty ; om:fromColumn a owl:ObjectProperty ; rdfs:label "from column" ; rdfs:subPropertyOf prov:used ; - rdfs:comment "Source column used in transformation" ; - rdfs:domain om:Column ; + rdfs:comment "Source column used in transformation. Object is an om:Column resource." ; + rdfs:domain om:ColumnLineage ; rdfs:range om:Column . om:toColumn a owl:ObjectProperty ; rdfs:label "to column" ; rdfs:subPropertyOf prov:wasGeneratedBy ; - rdfs:comment "Target column created by transformation" ; + rdfs:comment "Target column created by transformation. Object is an om:Column resource." ; + rdfs:domain om:ColumnLineage ; + rdfs:range om:Column . + +om:fromColumnFqn a owl:DatatypeProperty ; + rdfs:label "from column FQN" ; + rdfs:comment "FQN string of the source column. Retained alongside om:fromColumn for back-compat with consumers that match by string FQN." ; + rdfs:domain om:ColumnLineage ; + rdfs:range xsd:string . + +om:toColumnFqn a owl:DatatypeProperty ; + rdfs:label "to column FQN" ; + rdfs:comment "FQN string of the target column. Retained alongside om:toColumn for back-compat." ; + rdfs:domain om:ColumnLineage ; + rdfs:range xsd:string . + +om:hasChildColumn a owl:ObjectProperty ; + rdfs:label "has child column" ; + rdfs:comment "Nested child column (struct/map/union dataType). Subproperty of om:hasColumn." ; + rdfs:subPropertyOf om:hasColumn ; rdfs:domain om:Column ; rdfs:range om:Column . +om:hasColumnLineage a owl:ObjectProperty ; + rdfs:label "has column lineage" ; + rdfs:comment "Links a LineageDetails resource to a per-column lineage edge." ; + rdfs:domain om:LineageDetails ; + rdfs:range om:ColumnLineage . + +om:transformFunction a owl:DatatypeProperty ; + rdfs:label "transform function" ; + rdfs:comment "Function name applied between fromColumn(s) and toColumn (e.g. CONCAT, UPPER)." ; + rdfs:domain om:ColumnLineage ; + rdfs:range xsd:string . + om:lineageCreatedBy a owl:ObjectProperty ; rdfs:label "lineage created by" ; rdfs:subPropertyOf prov:wasAttributedTo ; diff --git a/openmetadata-spec/src/main/resources/rdf/shapes/openmetadata-shapes.ttl b/openmetadata-spec/src/main/resources/rdf/shapes/openmetadata-shapes.ttl index e8de71665729..d9060c48bef7 100644 --- a/openmetadata-spec/src/main/resources/rdf/shapes/openmetadata-shapes.ttl +++ b/openmetadata-spec/src/main/resources/rdf/shapes/openmetadata-shapes.ttl @@ -1,165 +1,213 @@ @prefix sh: . @prefix om: . @prefix xsd: . +@prefix rdf: . @prefix rdfs: . @prefix dcat: . -@prefix dcterms: . +@prefix dct: . +@prefix prov: . +@prefix skos: . -# SHACL Shapes for OpenMetadata Entities -# These shapes define validation constraints for RDF data +# SHACL shapes for the OpenMetadata knowledge graph. +# +# Predicate names here MUST match what the RDF mapper actually emits +# (see openmetadata-service/src/main/java/org/openmetadata/service/rdf/translator/). +# A shape that references a predicate the mapper never produces silently +# passes for every entity, which gives us false confidence. +# +# Cardinality is intentionally permissive: many fields are optional in +# OpenMetadata schemas. Shapes flag genuine inconsistencies (missing +# label/FQN, dangling references, type mismatches), not stylistic gaps. + +################################################################# +# Base Entity +################################################################# -# Base Entity Shape om:EntityShape a sh:NodeShape ; sh:targetClass om:Entity ; sh:property [ - sh:path om:id ; + sh:path rdfs:label ; sh:datatype xsd:string ; sh:minCount 1 ; - sh:maxCount 1 ; - sh:message "Entity must have exactly one ID"@en ; + sh:message "Every om:Entity must carry exactly one rdfs:label (its name)."@en ; ] ; sh:property [ - sh:path om:name ; + sh:path om:fullyQualifiedName ; sh:datatype xsd:string ; sh:minCount 1 ; sh:maxCount 1 ; - sh:pattern "^[a-zA-Z0-9_-]+$" ; - sh:message "Entity must have exactly one name with alphanumeric characters, hyphens, or underscores"@en ; + sh:message "Every om:Entity must have exactly one om:fullyQualifiedName."@en ; ] ; sh:property [ - sh:path om:fullyQualifiedName ; + sh:path dct:description ; sh:datatype xsd:string ; - sh:minCount 1 ; sh:maxCount 1 ; - sh:message "Entity must have exactly one fully qualified name"@en ; + sh:message "Description must be a single literal."@en ; ] ; sh:property [ - sh:path om:description ; - sh:datatype xsd:string ; + sh:path dct:hasVersion ; + sh:datatype xsd:decimal ; sh:maxCount 1 ; - sh:message "Entity can have at most one description"@en ; ] ; sh:property [ - sh:path om:version ; - sh:datatype xsd:decimal ; - sh:minInclusive 0.1 ; - sh:message "Version must be a positive decimal number"@en ; + sh:path om:hasOwner ; + sh:nodeKind sh:IRI ; + sh:message "om:hasOwner must reference a User or Team resource."@en ; ] . -# Table Shape +################################################################# +# Tables and Columns +################################################################# + om:TableShape a sh:NodeShape ; sh:targetClass om:Table ; sh:property [ sh:path om:hasColumn ; sh:class om:Column ; sh:minCount 1 ; - sh:message "Table must have at least one column"@en ; + sh:message "Table must declare at least one om:hasColumn link to an om:Column resource."@en ; ] ; sh:property [ - sh:path om:belongsTo ; - sh:class om:DatabaseSchema ; - sh:minCount 1 ; + sh:path om:belongsToSchema ; + sh:nodeKind sh:IRI ; sh:maxCount 1 ; - sh:message "Table must belong to exactly one database schema"@en ; ] ; sh:property [ - sh:path om:tableType ; - sh:datatype xsd:string ; - sh:in ("Regular" "External" "View" "SecureView" "MaterializedView" "Iceberg" "DatalakeTable") ; - sh:message "Table type must be one of the allowed values"@en ; + sh:path om:belongsToDatabase ; + sh:nodeKind sh:IRI ; + sh:maxCount 1 ; ] . -# Column Shape om:ColumnShape a sh:NodeShape ; sh:targetClass om:Column ; sh:property [ - sh:path om:dataType ; + sh:path om:fullyQualifiedName ; sh:datatype xsd:string ; sh:minCount 1 ; sh:maxCount 1 ; - sh:message "Column must have exactly one data type"@en ; + sh:message "Column must carry the FQN it was minted from."@en ; + ] ; + sh:property [ + sh:path om:columnDataType ; + sh:datatype xsd:string ; + sh:maxCount 1 ; ] ; sh:property [ sh:path om:ordinalPosition ; sh:datatype xsd:integer ; sh:minInclusive 0 ; - sh:message "Ordinal position must be a non-negative integer"@en ; ] ; sh:property [ - sh:path om:constraint ; - sh:datatype xsd:string ; - sh:in ("NULL" "NOT_NULL" "UNIQUE" "PRIMARY_KEY") ; - sh:message "Column constraint must be one of the allowed values"@en ; - ] . - -# Database Shape -om:DatabaseShape a sh:NodeShape ; - sh:targetClass om:Database ; + sh:path om:isPrimaryKey ; + sh:datatype xsd:boolean ; + sh:maxCount 1 ; + ] ; sh:property [ - sh:path om:hasSchema ; - sh:class om:DatabaseSchema ; - sh:message "Database can only contain database schemas"@en ; + sh:path om:isNullable ; + sh:datatype xsd:boolean ; + sh:maxCount 1 ; ] ; sh:property [ - sh:path om:belongsTo ; - sh:class om:DatabaseService ; - sh:minCount 1 ; + sh:path om:isUnique ; + sh:datatype xsd:boolean ; sh:maxCount 1 ; - sh:message "Database must belong to exactly one database service"@en ; + ] ; + sh:property [ + sh:path om:references ; + sh:class om:Column ; + sh:message "om:references must point at another om:Column resource (FK target)."@en ; ] . -# Database Schema Shape -om:DatabaseSchemaShape a sh:NodeShape ; - sh:targetClass om:DatabaseSchema ; +om:TableConstraintShape a sh:NodeShape ; + sh:targetClass om:TableConstraint ; + sh:property [ + sh:path om:constraintType ; + sh:datatype xsd:string ; + sh:minCount 1 ; + sh:maxCount 1 ; + sh:in ("PRIMARY_KEY" "UNIQUE" "FOREIGN_KEY" "SORT_KEY" "DIST_KEY" "CLUSTER_KEY") ; + sh:message "Constraint type must be one of the canonical TableConstraint enum values."@en ; + ] ; sh:property [ - sh:path om:hasTable ; - sh:class om:Table ; - sh:message "Database schema can only contain tables"@en ; + sh:path om:relationshipType ; + sh:datatype xsd:string ; + sh:in ("ONE_TO_ONE" "ONE_TO_MANY" "MANY_TO_ONE" "MANY_TO_MANY") ; + sh:maxCount 1 ; ] ; sh:property [ - sh:path om:belongsTo ; - sh:class om:Database ; + sh:path om:hasConstrainedColumn ; + sh:class om:Column ; sh:minCount 1 ; - sh:maxCount 1 ; - sh:message "Database schema must belong to exactly one database"@en ; + sh:message "A TableConstraint must reference at least one constrained Column."@en ; + ] ; + sh:property [ + sh:path om:hasReferredColumn ; + sh:class om:Column ; ] . -# Pipeline Shape -om:PipelineShape a sh:NodeShape ; - sh:targetClass om:Pipeline ; +################################################################# +# Lineage +################################################################# + +om:LineageDetailsShape a sh:NodeShape ; + sh:targetClass om:LineageDetails ; sh:property [ - sh:path om:hasTask ; - sh:class om:Task ; - sh:message "Pipeline can only contain tasks"@en ; + sh:path om:sqlQuery ; + sh:datatype xsd:string ; + sh:maxCount 1 ; ] ; sh:property [ - sh:path om:pipelineType ; + sh:path om:lineageSource ; sh:datatype xsd:string ; - sh:in ("airflow" "glue" "airbyte" "fivetran" "dagster" "nifi" "dbtCloud" "domoPipeline" "kafka" "kafkaConnect" "kinesis" "matillion" "domo" "flink" "databricksPipeline" "spline" "spark" "openLineage" "stitch" "dbt") ; - sh:message "Pipeline type must be one of the allowed values"@en ; + sh:maxCount 1 ; + ] ; + sh:property [ + sh:path prov:wasGeneratedBy ; + sh:nodeKind sh:IRI ; + sh:maxCount 1 ; + sh:message "If a pipeline produced this lineage, prov:wasGeneratedBy must reference exactly one resource."@en ; + ] ; + sh:property [ + sh:path om:hasColumnLineage ; + sh:class om:ColumnLineage ; ] . -# Dashboard Shape -om:DashboardShape a sh:NodeShape ; - sh:targetClass om:Dashboard ; +om:ColumnLineageShape a sh:NodeShape ; + sh:targetClass om:ColumnLineage ; + sh:property [ + sh:path om:fromColumn ; + sh:class om:Column ; + sh:message "om:fromColumn must reference an om:Column resource (URI), not an FQN literal — the FQN is carried by om:fromColumnFqn."@en ; + ] ; + sh:property [ + sh:path om:toColumn ; + sh:class om:Column ; + sh:maxCount 1 ; + sh:message "om:toColumn must reference exactly one om:Column resource."@en ; + ] ; + sh:property [ + sh:path om:fromColumnFqn ; + sh:datatype xsd:string ; + ] ; sh:property [ - sh:path om:hasChart ; - sh:class om:Chart ; - sh:message "Dashboard can only contain charts"@en ; + sh:path om:toColumnFqn ; + sh:datatype xsd:string ; + sh:maxCount 1 ; ] ; sh:property [ - sh:path om:dashboardType ; + sh:path om:transformFunction ; sh:datatype xsd:string ; - sh:in ("tableau" "looker" "superset" "redash" "metabase" "powerbi" "mode" "sigma" "lightdash" "mstr" "qlikSense" "quickSight" "domo" "databricks") ; - sh:message "Dashboard type must be one of the allowed values"@en ; + sh:maxCount 1 ; ] . -# Lineage Validation -om:LineageShape a sh:NodeShape ; +# A DataAsset must not appear as both upstream and downstream of the same +# other asset — that's a self-cycle through a single edge, almost always a +# bug in lineage capture. +om:LineageCycleShape a sh:NodeShape ; sh:targetClass om:DataAsset ; sh:sparql [ a sh:SPARQLConstraint ; - sh:message "Data asset cannot have cyclic lineage (be both upstream and downstream of the same asset)"@en ; + sh:message "Data asset is reported as both upstream and downstream of the same asset."@en ; sh:prefixes [ sh:declare [ sh:prefix "om" ; sh:namespace "https://open-metadata.org/ontology/"^^xsd:anyURI ; @@ -173,127 +221,113 @@ om:LineageShape a sh:NodeShape ; """ ; ] . -# Ownership Validation -om:OwnershipShape a sh:NodeShape ; - sh:targetClass om:Entity ; +################################################################# +# Containers +################################################################# + +om:DatabaseShape a sh:NodeShape ; + sh:targetClass om:Database ; sh:property [ - sh:path om:ownedBy ; - sh:or ( - [ sh:class om:User ] - [ sh:class om:Team ] - ) ; + sh:path om:belongsToService ; + sh:nodeKind sh:IRI ; sh:maxCount 1 ; - sh:message "Entity can be owned by at most one user or team"@en ; ] . -# Tag Application Shape -om:TagApplicationShape a sh:NodeShape ; - sh:targetClass om:Entity ; - sh:property [ - sh:path om:taggedWith ; - sh:class om:Tag ; - sh:message "Entity can only be tagged with valid tags"@en ; - ] ; +om:DatabaseSchemaShape a sh:NodeShape ; + sh:targetClass om:DatabaseSchema ; sh:property [ - sh:path om:classifiedAs ; - sh:class om:Classification ; - sh:message "Entity can only be classified with valid classifications"@en ; + sh:path om:belongsToDatabase ; + sh:nodeKind sh:IRI ; + sh:maxCount 1 ; ] . -# Domain Membership Shape -om:DomainMembershipShape a sh:NodeShape ; - sh:targetClass om:DataAsset ; +################################################################# +# Governance: Glossary, Domain, DataProduct, Tag +################################################################# + +om:GlossaryTermShape a sh:NodeShape ; + sh:targetClass om:GlossaryTerm ; sh:property [ - sh:path om:inDomain ; - sh:class om:Domain ; + sh:path skos:inScheme ; + sh:nodeKind sh:IRI ; + sh:minCount 1 ; sh:maxCount 1 ; - sh:message "Data asset can belong to at most one domain"@en ; - ] . - -# Data Quality Shape -om:DataQualityShape a sh:NodeShape ; - sh:targetClass om:Table ; + sh:message "Every glossary term must declare its parent glossary via skos:inScheme."@en ; + ] ; sh:property [ - sh:path om:dataQualityScore ; - sh:datatype xsd:decimal ; - sh:minInclusive 0.0 ; - sh:maxInclusive 100.0 ; - sh:message "Data quality score must be between 0 and 100"@en ; + sh:path skos:broader ; + sh:class om:GlossaryTerm ; + sh:maxCount 1 ; + sh:message "skos:broader must reference at most one parent term."@en ; ] ; sh:property [ - sh:path om:testCasePassed ; - sh:datatype xsd:integer ; - sh:minInclusive 0 ; - sh:message "Number of passed test cases must be non-negative"@en ; + sh:path skos:narrower ; + sh:class om:GlossaryTerm ; ] ; sh:property [ - sh:path om:testCaseFailed ; - sh:datatype xsd:integer ; - sh:minInclusive 0 ; - sh:message "Number of failed test cases must be non-negative"@en ; + sh:path skos:altLabel ; + sh:datatype xsd:string ; ] . -# Service Connection Shape -om:ServiceConnectionShape a sh:NodeShape ; - sh:targetClass om:Service ; +om:TagShape a sh:NodeShape ; + sh:targetClass om:Tag ; sh:property [ - sh:path om:serviceType ; - sh:datatype xsd:string ; - sh:minCount 1 ; + sh:path skos:inScheme ; + sh:nodeKind sh:IRI ; sh:maxCount 1 ; - sh:message "Service must have exactly one service type"@en ; + sh:message "Tag should declare its parent classification via skos:inScheme."@en ; ] ; sh:property [ - sh:path om:connectionConfig ; - sh:minCount 1 ; + sh:path om:tagFQN ; + sh:datatype xsd:string ; sh:maxCount 1 ; - sh:message "Service must have exactly one connection configuration"@en ; ] . -# Glossary Term Shape -om:GlossaryTermShape a sh:NodeShape ; - sh:targetClass om:GlossaryTerm ; +om:DomainShape a sh:NodeShape ; + sh:targetClass om:Domain ; sh:property [ - sh:path om:belongsTo ; - sh:class om:Glossary ; - sh:minCount 1 ; + sh:path om:domainType ; + sh:datatype xsd:string ; sh:maxCount 1 ; - sh:message "Glossary term must belong to exactly one glossary"@en ; ] ; sh:property [ - sh:path om:synonymOf ; - sh:class om:GlossaryTerm ; - sh:message "Synonym must reference another glossary term"@en ; - ] ; + sh:path om:hasDataProduct ; + sh:class om:DataProduct ; + ] . + +om:DataProductShape a sh:NodeShape ; + sh:targetClass om:DataProduct ; sh:property [ - sh:path om:antonymOf ; - sh:class om:GlossaryTerm ; - sh:message "Antonym must reference another glossary term"@en ; + sh:path om:hasAsset ; + sh:nodeKind sh:IRI ; ] . -# Team Hierarchy Shape -om:TeamHierarchyShape a sh:NodeShape ; - sh:targetClass om:Team ; - sh:sparql [ - a sh:SPARQLConstraint ; - sh:message "Team cannot be its own parent"@en ; - sh:prefixes [ sh:declare [ - sh:prefix "om" ; - sh:namespace "https://open-metadata.org/ontology/"^^xsd:anyURI ; - ] ] ; - sh:select """ - SELECT $this - WHERE { - $this om:isPartOf+ $this . - } - """ ; +################################################################# +# Quality +################################################################# + +om:TestCaseShape a sh:NodeShape ; + sh:targetClass om:TestCase ; + sh:property [ + sh:path om:appliesToEntity ; + sh:datatype xsd:string ; + sh:maxCount 1 ; + sh:message "TestCase should record the FQN of the entity it tests."@en ; + ] ; + sh:property [ + sh:path om:hasStatus ; + sh:datatype xsd:string ; + sh:maxCount 1 ; ] . -# User Team Membership Shape -om:UserTeamMembershipShape a sh:NodeShape ; - sh:targetClass om:User ; - sh:property [ - sh:path om:memberOf ; - sh:class om:Team ; - sh:message "User can only be member of valid teams"@en ; - ] . \ No newline at end of file +################################################################# +# Services +################################################################# + +om:ServiceShape a sh:NodeShape ; + sh:targetClass om:Service ; + sh:property [ + sh:path om:hasServiceType ; + sh:datatype xsd:string ; + sh:maxCount 1 ; + ] . diff --git a/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/box/box.tsx b/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/box/box.tsx index 7c1869a099ca..b4dd99be870e 100644 --- a/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/box/box.tsx +++ b/openmetadata-ui-core-components/src/main/resources/ui/src/components/base/box/box.tsx @@ -18,7 +18,7 @@ import { import { cx } from '@/utils/cx'; import type { HTMLAttributes, ReactNode } from 'react'; -export type Direction = 'row' | 'col' | 'row-reverse' | 'col-reverse'; +type Direction = 'row' | 'col' | 'row-reverse' | 'col-reverse'; type Align = 'start' | 'center' | 'end' | 'stretch' | 'baseline'; type Justify = 'start' | 'center' | 'end' | 'between' | 'around' | 'evenly'; type Wrap = 'wrap' | 'nowrap' | 'wrap-reverse'; diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ActivityAPI.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ActivityAPI.spec.ts index de9510aabe8c..569df2d01b66 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ActivityAPI.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ActivityAPI.spec.ts @@ -10,488 +10,907 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -import { expect } from '@playwright/test'; -import { DOMAIN_TAGS } from '../../constant/config'; -import { EntityTypeEndpoint } from '../../support/entity/Entity.interface'; +import { expect, Page, test as base } from '@playwright/test'; import { TableClass } from '../../support/entity/TableClass'; import { TagClass } from '../../support/tag/TagClass'; +import { AdminClass } from '../../support/user/AdminClass'; +import { UserClass } from '../../support/user/UserClass'; +import { createAdminApiContext, performAdminLogin } from '../../utils/admin'; import { - ACTIVITY_TEST_TIMEOUT, - addTagToTable, - createConversationThread, - createDescriptionActivityEventFromPage, - FEED_ITEM_TIMEOUT, - getActivityFeedItems, - getFeedItemByText, - getTableFqn, - getTableLeafName, - openActivityFeedAndWaitForApi, - patchTableDescription, - THUMBS_UP_EMOJI, - toggleThumbsUpReaction, - visitTableActivityFeed, - waitForActivityEvent, -} from '../../utils/activityAPI'; -import { postActivityComment } from '../../utils/activityFeed'; -import { performAdminLogin } from '../../utils/admin'; -import { getApiContext, redirectToHomePage, uuid } from '../../utils/common'; -import { - addOwner, - updateDescription, - waitForAllLoadersToDisappear, -} from '../../utils/entity'; -import { test } from '../fixtures/pages'; - -test.describe( - 'Activity API - Entity Changes', - { tag: [DOMAIN_TAGS.DISCOVERY] }, - () => { - let entityChangesTable: TableClass; - let entityChangesTag: TagClass; - let adminDisplayName: string; - - test.beforeAll('Setup: create table and tag', async ({ browser }) => { - const { apiContext, afterAction } = await performAdminLogin(browser); + descriptionBox, + getApiContext, + redirectToHomePage, + uuid, +} from '../../utils/common'; +import { waitForPageLoaded } from '../../utils/polling'; + +let adminUser: AdminClass; +let testTable: TableClass; +let testTag: TagClass; + +const test = base.extend<{ + page: Page; +}>({ + page: async ({ browser }, use) => { + const adminPage = await browser.newPage(); + await adminUser.login( + adminPage, + adminUser.data.email, + adminUser.data.password + ); + await use(adminPage); + await adminPage.close(); + }, +}); + +type ActivityApiEvent = { + actor?: { displayName?: string; name?: string }; + eventType?: string; + summary?: string; +}; + +type ActivityApiResponse = { + data?: ActivityApiEvent[]; +}; + +const openActivityFeedAndWaitForApi = async (page: Page, entityFqn: string) => { + const expectedActivityPath = `/api/v1/activity/entity/table/name/${entityFqn}`; + const activityResponsePromise = page.waitForResponse( + (response) => + decodeURIComponent(response.url()).includes(expectedActivityPath) && + response.status() === 200, + { timeout: 15000 } + ); + + await page.getByTestId('activity_feed').click(); + await waitForPageLoaded(page); + + const activityResponse = await activityResponsePromise; + + return (await activityResponse.json()) as ActivityApiResponse; +}; + +const waitForActivityEvent = async (entityFqn: string, eventType: string) => { + const { apiContext, afterAction } = await createAdminApiContext(); + const activityUrl = `/api/v1/activity/entity/table/name/${encodeURIComponent( + entityFqn + )}?days=30&limit=50`; + let events: ActivityApiEvent[] = []; + + try { + await expect + .poll( + async () => { + const response = await apiContext.get(activityUrl); + + if (!response.ok()) { + return false; + } + + const body = (await response.json()) as ActivityApiResponse; + events = body.data ?? []; + + return events.some((event) => event.eventType === eventType); + }, + { + timeout: 300000, // 5 minutes + intervals: [1000, 2000, 5000, 10000], + message: `Timed out waiting for ${eventType} event for ${entityFqn}`, + } + ) + .toBe(true); - entityChangesTable = new TableClass(); - entityChangesTag = new TagClass({}); + return events.find((event) => event.eventType === eventType); + } finally { + await afterAction(); + } +}; - try { - await entityChangesTable.create(apiContext); - await entityChangesTag.create(apiContext); +const addOwnerFromActivitySpec = async (page: Page, owner: string) => { + await page.getByTestId('edit-owner').click(); - const userResponse = await apiContext.get('/api/v1/users/loggedInUser'); - const adminUser = await userResponse.json(); - adminDisplayName = adminUser.displayName ?? adminUser.name; - } finally { - await afterAction(); - } - }); + const usersTab = page.getByRole('tab', { name: /^Users\b/ }); + if ((await usersTab.getAttribute('aria-selected')) !== 'true') { + await usersTab.click(); + } - test.beforeEach(async ({ page }) => { - await redirectToHomePage(page); - await waitForAllLoadersToDisappear(page); - }); + const ownerSearchInput = page.getByTestId('owner-select-users-search-bar'); + await expect(ownerSearchInput).toBeVisible({ timeout: 30000 }); - test('creates an activity event when the description is updated', async ({ - page, - }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); - - const newDescription = `Test description updated at ${Date.now()}`; - const entityFqn = getTableFqn(entityChangesTable); - - await test.step('Update the table description from the entity page', async () => { - await entityChangesTable.visitEntityPage(page); - await waitForAllLoadersToDisappear(page); - await updateDescription( - page, - newDescription, - false, - 'asset-description-container', - EntityTypeEndpoint.Table - ); - }); + const searchUser = page.waitForResponse( + (response) => + response.url().includes('/api/v1/search/query') && response.ok() + ); + await ownerSearchInput.fill(owner); + await searchUser; - await test.step('Verify the description event through API and UI', async () => { - const descriptionEvent = await waitForActivityEvent({ - entityFqn, - eventType: 'DescriptionUpdated', - text: newDescription, - }); - const activityResponse = await openActivityFeedAndWaitForApi( - page, - entityFqn - ); - const renderedDescriptionEvent = activityResponse.data?.find( - (event) => - event.eventType === 'DescriptionUpdated' && - JSON.stringify(event).includes(newDescription) - ); - const feedItem = await getFeedItemByText(page, newDescription); + const ownerItem = page.getByRole('listitem', { name: owner }); + await expect(ownerItem).toBeVisible({ timeout: 60000 }); + await ownerItem.click(); - expect(descriptionEvent).toBeDefined(); - expect(renderedDescriptionEvent).toBeDefined(); - await expect(feedItem).toContainText(/description/i); - }); - }); + const patchRequest = page.waitForResponse('/api/v1/tables/*'); + await page + .locator('[id^="rc-tabs-"][id$="-panel-users"]') + .getByTestId('selectable-list-update-btn') + .click(); + await patchRequest; - test('creates an activity event when tags are added', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); + await expect(page.getByTestId('owner-link').getByTestId(owner)).toBeVisible(); +}; - const entityFqn = getTableFqn(entityChangesTable); - const tagDisplayName = entityChangesTag.getTagDisplayName(); +test.beforeAll('Setup create admin user', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); + adminUser = new AdminClass(); - await test.step('Add a tag to the table through API setup', async () => { - const { apiContext, afterAction } = await getApiContext(page); + await adminUser.create(apiContext); - try { - await addTagToTable(apiContext, entityChangesTable, entityChangesTag); - } finally { - await afterAction(); - } - }); + await afterAction(); +}); - await test.step('Verify the tag event through API and UI', async () => { - const tagsEvent = await waitForActivityEvent({ - entityFqn, - eventType: 'TagsUpdated', - }); - const activityResponse = await visitTableActivityFeed( - page, - entityChangesTable - ); - const renderedTagsEvent = activityResponse.data?.find( - (event) => event.eventType === 'TagsUpdated' - ); - const feedItem = getActivityFeedItems(page) - .filter({ hasText: /tag/i }) - .filter({ hasText: tagDisplayName }); +test.afterAll('Cleanup delete admin user', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); - expect(tagsEvent).toBeDefined(); - expect(renderedTagsEvent).toBeDefined(); - await expect(feedItem).toBeVisible({ timeout: FEED_ITEM_TIMEOUT }); - }); - }); - - test('creates an activity event when owner is added', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); - - const entityFqn = getTableFqn(entityChangesTable); - - await test.step('Add the owner from the entity page', async () => { - await entityChangesTable.visitEntityPage(page); - await waitForAllLoadersToDisappear(page); - await addOwner({ - page, - owner: adminDisplayName, - endpoint: EntityTypeEndpoint.Table, - }); - }); + await adminUser.delete(apiContext); - await test.step('Verify the owner event through API and UI', async () => { - const ownerEvent = await waitForActivityEvent({ - entityFqn, - eventType: 'OwnerUpdated', - }); - const activityResponse = await openActivityFeedAndWaitForApi( - page, - entityFqn - ); - const renderedOwnerEvent = activityResponse.data?.find( - (event) => event.eventType === 'OwnerUpdated' - ); - const feedItem = getActivityFeedItems(page) - .filter({ hasText: /owner/i }) - .filter({ hasText: adminDisplayName }); + await afterAction(); +}); - expect(ownerEvent).toBeDefined(); - expect(renderedOwnerEvent).toBeDefined(); - await expect(feedItem).toBeVisible({ timeout: FEED_ITEM_TIMEOUT }); - }); - }); - - test('shows the actor who made the activity change', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); - - const entityFqn = getTableFqn(entityChangesTable); - const uniqueDescription = `Actor test description ${Date.now()}`; - - await test.step('Make a table change as the logged-in admin user', async () => { - const { apiContext, afterAction } = await getApiContext(page); - - try { - await patchTableDescription( - apiContext, - entityChangesTable, - uniqueDescription - ); - } finally { - await afterAction(); - } +test.describe('Activity API - Entity Changes', () => { + test.beforeAll('Setup: create entities and users', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); + testTable = new TableClass(); + testTag = new TagClass({}); - await waitForActivityEvent({ - entityFqn, - eventType: 'DescriptionUpdated', - text: uniqueDescription, - }); - }); + await testTable.create(apiContext); + await testTag.create(apiContext); - await test.step('Verify the actor is visible in the matching feed item', async () => { - await visitTableActivityFeed(page, entityChangesTable); + await afterAction(); + }); - const feedItem = await getFeedItemByText(page, uniqueDescription); + test.afterAll('Cleanup: delete entities and users', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); - await expect(feedItem).toContainText(adminDisplayName); - }); - }); + await testTable.delete(apiContext); + await testTag.delete(apiContext); - test('links activity items to the correct entity', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); + await afterAction(); + }); - const description = `Entity link description ${uuid()}`; + test.beforeEach(async ({ page }) => { + await redirectToHomePage(page); + await waitForPageLoaded(page); + }); - await test.step('Create an activity event for the table', async () => { - await createDescriptionActivityEventFromPage( - page, - entityChangesTable, - description - ); - }); + test('Activity event is created when description is updated', async ({ + page, + }) => { + test.setTimeout(300000); + const newDescription = `Test description updated at ${Date.now()}`; + const entityFqn = testTable.entityResponseData.fullyQualifiedName ?? ''; - await test.step('Verify the feed card has the table entity link', async () => { - await visitTableActivityFeed(page, entityChangesTable); + // Navigate to entity page + await testTable.visitEntityPage(page); - const feedItem = await getFeedItemByText(page, description); - const entityLink = feedItem.locator('a[href*="/table/"]').first(); + // Update description + await page.getByTestId('edit-description').click(); - await expect(entityLink).toBeVisible(); + // Wait for description modal to appear + await page.locator(descriptionBox).waitFor({ state: 'visible' }); + await page.locator(descriptionBox).clear(); + await page.locator(descriptionBox).fill(newDescription); - const href = await entityLink.getAttribute('href'); + const patchResponse = page.waitForResponse('/api/v1/tables/*'); + await page.getByTestId('save').click(); + await patchResponse; - expect(href).toContain('table'); - expect(href).toContain(getTableLeafName(entityChangesTable)); - }); - }); - } -); + const descriptionEvent = await waitForActivityEvent( + entityFqn, + 'DescriptionUpdated' + ); + const activityResponse = await openActivityFeedAndWaitForApi( + page, + entityFqn + ); + const feedContainer = page.locator( + '#center-container [data-testid="message-container"]' + ); + const renderedDescriptionEvent = activityResponse.data?.find( + (event) => event.eventType === 'DescriptionUpdated' + ); + + expect(descriptionEvent).toBeDefined(); + expect(renderedDescriptionEvent).toBeDefined(); + await expect(feedContainer.first()).toContainText(/description/i); + }); + + test('Activity event is created when tags are added', async ({ page }) => { + test.setTimeout(300000); + const entityFqn = testTable.entityResponseData.fullyQualifiedName ?? ''; + + // Add tag via API to bypass search indexing issues + const { apiContext, afterAction } = await getApiContext(page); + + // Add tag to the table via API + await apiContext.patch( + `/api/v1/tables/${testTable.entityResponseData.id}`, + { + data: [ + { + op: 'add', + path: '/tags/0', + value: { + tagFQN: testTag.responseData.fullyQualifiedName, + source: 'Classification', + }, + }, + ], + headers: { + 'Content-Type': 'application/json-patch+json', + }, + } + ); -test.describe( - 'Activity API - Reactions', - { tag: [DOMAIN_TAGS.DISCOVERY] }, - () => { - const reactionsTable = new TableClass(); + await afterAction(); - test.beforeAll('Setup: create table', async ({ browser }) => { - const { apiContext, afterAction } = await performAdminLogin(browser); + const tagsEvent = await waitForActivityEvent(entityFqn, 'TagsUpdated'); - try { - await reactionsTable.create(apiContext); - } finally { - await afterAction(); + // Navigate to entity page + await testTable.visitEntityPage(page); + + const activityResponse = await openActivityFeedAndWaitForApi( + page, + entityFqn + ); + const feedContainer = page.locator( + '#center-container [data-testid="message-container"]' + ); + const renderedTagsEvent = activityResponse.data?.find( + (event) => event.eventType === 'TagsUpdated' + ); + + expect(tagsEvent).toBeDefined(); + expect(renderedTagsEvent).toBeDefined(); + await expect(feedContainer.first()).toContainText(/tag/i); + }); + + test('Activity event is created when owner is added', async ({ page }) => { + test.setTimeout(300000); + const entityFqn = testTable.entityResponseData.fullyQualifiedName ?? ''; + const ownerDisplayName = adminUser.getUserDisplayName(); + + // Navigate to entity page + await testTable.visitEntityPage(page); + + await addOwnerFromActivitySpec(page, ownerDisplayName); + + const ownerEvent = await waitForActivityEvent(entityFqn, 'OwnerUpdated'); + const activityResponse = await openActivityFeedAndWaitForApi( + page, + entityFqn + ); + const feedContainer = page.locator( + '#center-container [data-testid="message-container"]' + ); + const renderedOwnerEvent = activityResponse.data?.find( + (event) => event.eventType === 'OwnerUpdated' + ); + + expect(ownerEvent).toBeDefined(); + expect(renderedOwnerEvent).toBeDefined(); + await expect(feedContainer.first()).toContainText(/owner/i); + await expect(feedContainer.first()).toContainText(ownerDisplayName); + }); + + test('Activity event shows the actor who made the change', async ({ + page, + }) => { + test.setTimeout(300000); + // Make a change via API so we know exactly who the actor is + const { apiContext, afterAction } = await getApiContext(page); + const uniqueDescription = `Actor test description ${Date.now()}`; + + await apiContext.patch( + `/api/v1/tables/${testTable.entityResponseData.id}`, + { + data: [ + { + op: 'add', + path: '/description', + value: uniqueDescription, + }, + ], + headers: { + 'Content-Type': 'application/json-patch+json', + }, } - }); + ); - test.beforeEach(async ({ page }) => { - await redirectToHomePage(page); - await waitForAllLoadersToDisappear(page); - }); + await afterAction(); - test('adds a reaction to a feed item', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); + // Wait for activity to be indexed + await waitForActivityEvent( + testTable.entityResponseData.fullyQualifiedName ?? '', + 'DescriptionUpdated' + ); - const description = `Test activity for adding reaction ${uuid()}`; + // Navigate to entity page + await testTable.visitEntityPage(page); - await test.step('Create and open an activity feed item', async () => { - await createDescriptionActivityEventFromPage( - page, - reactionsTable, - description - ); - await visitTableActivityFeed(page, reactionsTable); + // Navigate to Activity Feed tab + await page.getByTestId('activity_feed').click(); + await waitForPageLoaded(page); + + // Check if there are any feed items + const feedContainer = page + .locator('#center-container [data-testid="message-container"]') + .filter({ + hasText: uniqueDescription, }); - await test.step('Add thumbs-up reaction and verify it is visible', async () => { - const feedItem = await getFeedItemByText(page, description); + await feedContainer.waitFor({ state: 'visible' }); - await toggleThumbsUpReaction(feedItem, page); - await expect( - feedItem.getByRole('button', { name: new RegExp(THUMBS_UP_EMOJI) }) - ).toBeVisible({ timeout: 5_000 }); - }); - }); + const feedContent = await feedContainer.first().textContent(); - test('removes an existing reaction from a feed item', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); + // The activity should show the actor's name (admin user who made the change) + // Activity typically shows username or display name + const actorName = adminUser.responseData.displayName; - const description = `Test activity for removing reaction ${uuid()}`; + expect(feedContent).toContain(actorName); + }); - await test.step('Create and open an activity feed item', async () => { - await createDescriptionActivityEventFromPage( - page, - reactionsTable, - description - ); - await visitTableActivityFeed(page, reactionsTable); - }); + test('Activity event links to the correct entity', async ({ page }) => { + // Navigate to entity page + await testTable.visitEntityPage(page); + + // Navigate to Activity Feed tab + await page.getByTestId('activity_feed').click(); + await waitForPageLoaded(page); + + // Check if there are any feed items + const feedContainer = page.locator( + '#center-container [data-testid="message-container"]' + ); - await test.step('Add and then remove thumbs-up reaction', async () => { - const feedItem = await getFeedItemByText(page, description); + await feedContainer + .first() + .waitFor({ state: 'visible', timeout: 10000 }) + .catch(() => {}); - await toggleThumbsUpReaction(feedItem, page); - await expect( - feedItem.getByRole('button', { name: new RegExp(THUMBS_UP_EMOJI) }) - ).toBeVisible({ timeout: 5_000 }); + if ((await feedContainer.count()) > 0) { + // Activity cards now render actor/profile links ahead of the entity link. + const entityLink = feedContainer + .first() + .locator('a[href*="/table/"]') + .first(); + + if (await entityLink.isVisible()) { + const href = await entityLink.getAttribute('href'); - await toggleThumbsUpReaction(feedItem, page); - await expect( - feedItem.getByRole('button', { name: new RegExp(THUMBS_UP_EMOJI) }) - ).not.toBeVisible({ timeout: 5_000 }); + // The link should point to the table entity + expect(href).toContain('table'); + expect(href).toContain( + (testTable.entityResponseData.fullyQualifiedName ?? '') + .split('.') + .pop() ?? '' + ); + } else { + // Some activity types may not have clickable entity links + test.info().annotations.push({ + type: 'note', + description: 'No entity link found in activity item', + }); + } + } else { + test.info().annotations.push({ + type: 'note', + description: 'Activity feed is empty - cannot verify entity link', }); - }); - } -); + } + }); +}); -test.describe( - 'Activity API - Comments', - { tag: [DOMAIN_TAGS.DISCOVERY] }, - () => { - const commentsTable = new TableClass(); +test.describe('Activity API - Reactions', () => { + const adminUser = new UserClass(); + const testTable = new TableClass(); - test.beforeAll('Setup: create table', async ({ browser }) => { + test.beforeAll( + 'Setup: create entities and conversation', + async ({ browser }) => { const { apiContext, afterAction } = await performAdminLogin(browser); try { - await commentsTable.create(apiContext); + await adminUser.create(apiContext); + await adminUser.setAdminRole(apiContext); + await testTable.create(apiContext); + + // Create a conversation thread to have something to react to + const entityLink = `<#E::table::${testTable.entityResponseData.fullyQualifiedName}>`; + await apiContext.post('/api/v1/feed', { + data: { + message: 'Test conversation for reactions', + about: entityLink, + }, + }); } finally { await afterAction(); } - }); - - test.beforeEach(async ({ page }) => { - await redirectToHomePage(page); - await waitForAllLoadersToDisappear(page); - }); - - test('adds a comment to a feed item', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); - - const description = `Test activity for comments ${uuid()}`; - const commentText = `Test comment ${uuid()}`; - - await test.step('Create and open an activity feed item', async () => { - await createDescriptionActivityEventFromPage( - page, - commentsTable, - description - ); - await visitTableActivityFeed(page, commentsTable); + } + ); + + test.afterAll('Cleanup: delete entities', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); + + try { + await testTable.delete(apiContext); + await adminUser.delete(apiContext); + } finally { + await afterAction(); + } + }); + + test.beforeEach(async ({ page }) => { + await redirectToHomePage(page); + await waitForPageLoaded(page); + }); + + test('User can add reaction to feed item', async ({ page }) => { + // Navigate to entity page + await testTable.visitEntityPage(page); + + // Navigate to Activity Feed tab + await page.getByTestId('activity_feed').click(); + await waitForPageLoaded(page); + + // Wait for feed to load + const feedContainer = page.locator( + '#center-container [data-testid="message-container"]' + ); + const emptyState = page.locator( + '[data-testid="no-data-placeholder-container"]' + ); + + // Wait for either feed items or empty state + await Promise.race([ + feedContainer.first().waitFor({ state: 'visible', timeout: 10000 }), + emptyState.waitFor({ state: 'visible', timeout: 10000 }), + ]).catch(() => {}); + + // Skip if no feed items + if ((await feedContainer.count()) === 0) { + test.info().annotations.push({ + type: 'skip', + description: 'No feed items available to react to', }); - await test.step('Open the feed detail and post a comment', async () => { - const feedItem = await getFeedItemByText(page, description); - - await feedItem.click(); - await waitForAllLoadersToDisappear(page); - await postActivityComment(page, commentText); + return; + } + + // Find reaction button and click + const reactionContainer = feedContainer + .first() + .locator('[data-testid="feed-reaction-container"]'); + + if (await reactionContainer.isVisible()) { + const addReactionBtn = reactionContainer.locator( + '[data-testid="add-reactions"]' + ); + + await expect(addReactionBtn).toBeVisible(); + await addReactionBtn.click(); + + // Wait for reaction popover + await page + .locator('.ant-popover-feed-reactions') + .waitFor({ state: 'visible' }); + + // Click on thumbsUp reaction + const reactionResponse = page.waitForResponse( + (response) => + (response.url().includes('/api/v1/feed') || + response.url().includes('/api/v1/activity')) && + response.status() === 200 + ); + + await page + .locator('[data-testid="reaction-button"][title="thumbsUp"]') + .click(); + + await reactionResponse; + + // Verify reaction is added - look for thumbsUp emoji in the feed + // The reaction button shows as "👍 1" or similar + const thumbsUpReaction = page.getByRole('button', { name: /👍/ }); + + await expect(thumbsUpReaction.first()).toBeVisible({ timeout: 5000 }); + } + }); + + test('User can remove reaction from feed item', async ({ page }) => { + // Navigate to entity page + await testTable.visitEntityPage(page); + + // Navigate to Activity Feed tab + await page.getByTestId('activity_feed').click(); + await waitForPageLoaded(page); + + // Wait for feed to load + const feedContainer = page.locator( + '#center-container [data-testid="message-container"]' + ); + const emptyState = page.locator( + '[data-testid="no-data-placeholder-container"]' + ); + + // Wait for either feed items or empty state + await Promise.race([ + feedContainer.first().waitFor({ state: 'visible', timeout: 10000 }), + emptyState.waitFor({ state: 'visible', timeout: 10000 }), + ]).catch(() => {}); + + // Skip if no feed items + if ((await feedContainer.count()) === 0) { + test.info().annotations.push({ + type: 'skip', + description: 'No feed items available', }); - }); - test('shows the activity detail layout', async ({ page }) => { - // waitForActivityEvent polls the API for up to 5 minutes (ACTIVITY_EVENT_TIMEOUT = 300_000ms) - // because activity events are processed asynchronously in the background. - test.setTimeout(ACTIVITY_TEST_TIMEOUT); + return; + } + + // Find reaction container + const reactionContainer = feedContainer + .first() + .locator('[data-testid="feed-reaction-container"]'); - const description = `Test activity detail layout ${uuid()}`; + if (await reactionContainer.isVisible()) { + // First add a reaction if not already present + const existingReaction = reactionContainer.locator( + '[data-testid="Reactions"]' + ); - await test.step('Create and open an activity feed item', async () => { - await createDescriptionActivityEventFromPage( - page, - commentsTable, - description + if (!(await existingReaction.isVisible())) { + const addReactionBtn = reactionContainer.locator( + '[data-testid="add-reactions"]' ); - await visitTableActivityFeed(page, commentsTable); + await addReactionBtn.click(); + await page + .locator('.ant-popover-feed-reactions') + .waitFor({ state: 'visible' }); + + const addResponse = page.waitForResponse( + (response) => + (response.url().includes('/api/v1/feed') || + response.url().includes('/api/v1/activity')) && + response.status() === 200 + ); + await page + .locator('[data-testid="reaction-button"][title="thumbsUp"]') + .click(); + await addResponse; + } + + // Now remove the reaction by clicking again + await reactionContainer.locator('[data-testid="add-reactions"]').click(); + await page + .locator('.ant-popover-feed-reactions') + .waitFor({ state: 'visible' }); + + const removeResponse = page.waitForResponse( + (response) => + (response.url().includes('/api/v1/feed') || + response.url().includes('/api/v1/activity')) && + response.status() === 200 + ); + + await page + .locator('[data-testid="reaction-button"][title="thumbsUp"]') + .click(); + + await removeResponse; + } + }); +}); + +test.describe('Activity API - Comments', () => { + let adminUser: UserClass; + let testTable: TableClass; + let feedMessage = ''; + let feedUrl = ''; + let entityLink = ''; + + test.beforeAll( + 'Setup: create entities and conversation', + async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); + + adminUser = new UserClass(); + testTable = new TableClass(); + + await adminUser.create(apiContext); + await adminUser.setAdminRole(apiContext); + await testTable.create(apiContext); + + // Create a conversation thread to have something to comment on + feedMessage = `Test conversation for comments ${Date.now()}`; + entityLink = `<#E::table::${testTable.entityResponseData.fullyQualifiedName}>`; + await apiContext.post('/api/v1/feed', { + data: { + message: feedMessage, + about: entityLink, + }, + }); + feedUrl = `/api/v1/feed?entityLink=${encodeURIComponent( + entityLink + )}&type=Conversation&limit=25`; + + await expect + .poll( + async () => { + const response = await apiContext.get(feedUrl); + const data = await response.json(); + + return (data.data ?? []).some((thread: { message?: string }) => + thread.message?.includes(feedMessage) + ); + }, + { timeout: 60_000, intervals: [2_000] } + ) + .toBe(true); + + await afterAction(); + } + ); + + test.afterAll('Cleanup: delete entities', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); + + try { + await testTable.delete(apiContext); + await adminUser.delete(apiContext); + } finally { + await afterAction(); + } + }); + + test.beforeEach(async ({ page }) => { + await redirectToHomePage(page); + await waitForPageLoaded(page); + }); + + test('User can add comment to feed item', async ({ page }) => { + const commentText = `Test comment ${uuid()}`; + + // Navigate to entity page + await testTable.visitEntityPage(page); + + // Navigate to Activity Feed tab + await page.getByTestId('activity_feed').click(); + await waitForPageLoaded(page); + + // Wait for feed to load + const feedContainer = page + .locator('#center-container [data-testid="message-container"]') + .filter({ + hasText: feedMessage, }); - await test.step('Open the detail view and verify layout regions', async () => { - const feedItem = await getFeedItemByText(page, description); + // Wait for either feed items or empty state - await feedItem.click(); - await waitForAllLoadersToDisappear(page); + feedContainer.waitFor({ state: 'visible' }); - const activityPanel = page.locator('#activity-panel'); + // Click on the feed card to open detail view + await feedContainer.click(); + await waitForPageLoaded(page); - await expect(activityPanel).toBeVisible(); - await expect( - activityPanel.getByTestId('comments-input-field') - ).toBeVisible(); - }); - }); - } -); + // Wait for comment input to appear + const commentInput = page.locator('[data-testid="comments-input-field"]'); -test.describe( - 'Activity API - Homepage Widget', - { tag: [DOMAIN_TAGS.DISCOVERY] }, - () => { - const homepageTable = new TableClass(); + if (await commentInput.isVisible()) { + await commentInput.click(); - test.beforeAll('Setup: create table and activity', async ({ browser }) => { - const { apiContext, afterAction } = await performAdminLogin(browser); + // Fill in the comment using the editor + const editorField = page.locator( + '[data-testid="editor-wrapper"] [contenteditable="true"]' + ); + await editorField.fill(commentText); - try { - await homepageTable.create(apiContext); - await createConversationThread( - apiContext, - homepageTable, - `Test conversation for homepage widget ${uuid()}` - ); - } finally { - await afterAction(); - } - }); + // Wait for send button to be enabled + const sendButton = page.getByTestId('send-button'); - test.beforeEach(async ({ page }) => { - await redirectToHomePage(page); - await waitForAllLoadersToDisappear(page); - }); + await expect(sendButton).toBeEnabled(); - test('displays feed content in the Activity Feed widget', async ({ - page, - }) => { - const feedWidget = page.getByTestId('KnowledgePanel.ActivityFeed'); - const feedItems = feedWidget.getByTestId('message-container'); + // Send the comment + const postResponse = page.waitForResponse('/api/v1/feed/*/posts'); + await sendButton.click(); + await postResponse; + + // Verify comment appears + await expect(page.getByText(commentText)).toBeVisible({ timeout: 10000 }); + } else { + test.info().annotations.push({ + type: 'note', + description: 'Comment input not visible for this feed item type', + }); + } + }); + + test('Feed detail shows conversation layout', async ({ page }) => { + // Navigate to entity page + await testTable.visitEntityPage(page); + + // Navigate to Activity Feed tab + await page.getByTestId('activity_feed').click(); + await waitForPageLoaded(page); + + // Wait for feed to load + const feedContainer = page.locator( + '#center-container [data-testid="message-container"]' + ); + const emptyState = page.locator( + '[data-testid="no-data-placeholder-container"]' + ); + + // Wait for either feed items or empty state + await Promise.race([ + feedContainer.first().waitFor({ state: 'visible', timeout: 10000 }), + emptyState.waitFor({ state: 'visible', timeout: 10000 }), + ]).catch(() => {}); + + // Skip if no feed items + if ((await feedContainer.count()) === 0) { + test.info().annotations.push({ + type: 'skip', + description: 'No feed items available', + }); - await expect(feedWidget).toBeVisible(); - await expect(feedItems.first()).toBeVisible({ - timeout: FEED_ITEM_TIMEOUT, + return; + } + + // Click on feed card to open detail view + await feedContainer.first().click(); + await waitForPageLoaded(page); + + // Verify layout elements are visible + // Feed card sidebar (at least one should be visible) + const leftSidebar = page + .locator('[data-testid="feed-card-v2-sidebar"]') + .first(); + + // Feed replies section + const feedReplies = page.locator('[data-testid="feed-replies"]'); + + // Feed content panel or activity feed panel + const feedContentPanel = page.locator('.activity-feed-content-panel'); + const activityFeedPanel = page.locator( + '[data-testid="activity-feed-panel"]' + ); + + // Check at least one layout element is visible + const hasLayout = + (await leftSidebar.isVisible().catch(() => false)) || + (await feedReplies.isVisible().catch(() => false)) || + (await feedContentPanel.isVisible().catch(() => false)) || + (await activityFeedPanel.isVisible().catch(() => false)); + + expect(hasLayout).toBe(true); + }); +}); + +test.describe('Activity API - Homepage Widget', () => { + const adminUser = new UserClass(); + const testTable = new TableClass(); + + test.beforeAll('Setup: create entities and activity', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); + + try { + await adminUser.create(apiContext); + await adminUser.setAdminRole(apiContext); + await testTable.create(apiContext); + + // Create a conversation to ensure there's activity in the feed + const entityLink = `<#E::table::${testTable.entityResponseData.fullyQualifiedName}>`; + await apiContext.post('/api/v1/feed', { + data: { + message: 'Test conversation for homepage widget', + about: entityLink, + }, + }); + } finally { + await afterAction(); + } + }); + + test.afterAll('Cleanup: delete entities', async ({ browser }) => { + const { apiContext, afterAction } = await performAdminLogin(browser); + + try { + await testTable.delete(apiContext); + await adminUser.delete(apiContext); + } finally { + await afterAction(); + } + }); + + test.beforeEach(async ({ page }) => { + await redirectToHomePage(page); + await waitForPageLoaded(page); + }); + + test('Activity Feed widget displays feed items', async ({ page }) => { + // Wait for feed widget to load + page + .getByTestId('KnowledgePanel.ActivityFeed') + .waitFor({ state: 'visible' }); + + // Check for feed content - either messages, empty state, or "No Recent Activity" + const messageContainers = page.locator( + '#center-container [data-testid="message-container"]' + ); + const noRecentActivity = page.getByText('No Recent Activity'); + const emptyState = page.locator( + '[data-testid="no-data-placeholder-container"]' + ); + + // Either we have messages or empty state + const hasMessages = (await messageContainers.count()) > 0; + const hasNoRecentActivity = await noRecentActivity + .isVisible() + .catch(() => false); + const hasEmpty = (await emptyState.count()) > 0; + + expect(hasMessages || hasNoRecentActivity || hasEmpty).toBe(true); + }); + + test('Activity Feed widget has filter options', async ({ page }) => { + // Wait for feed widget to load + const feedWidget = page.getByTestId('KnowledgePanel.ActivityFeed'); + + // Widget may not exist if homepage is not customized + const widgetExists = await feedWidget.isVisible().catch(() => false); + + if (!widgetExists) { + test.info().annotations.push({ + type: 'skip', + description: 'Activity Feed widget not visible on homepage', }); - }); - test('shows Activity Feed widget filter options', async ({ page }) => { - const feedWidget = page.getByTestId('KnowledgePanel.ActivityFeed'); + return; + } - await expect(feedWidget).toBeVisible(); + await expect(feedWidget).toBeVisible(); - const sortDropdown = feedWidget.getByTestId('widget-sort-by-dropdown'); + // Find sort dropdown in widget header + const sortDropdown = feedWidget.getByTestId('widget-sort-by-dropdown'); - await expect(sortDropdown).toBeVisible(); - await expect(sortDropdown).toBeEnabled(); - await sortDropdown.click(); + // Sort dropdown may not be visible if widget is in minimal mode + const dropdownExists = await sortDropdown.isVisible().catch(() => false); - const filterMenu = page.getByRole('menu').filter({ - hasText: 'All Activity', + if (!dropdownExists) { + test.info().annotations.push({ + type: 'note', + description: 'Sort dropdown not visible in widget header', }); - await expect(filterMenu).toBeVisible(); - await expect( - page.getByRole('menuitem', { name: 'All Activity' }) - ).toBeVisible(); - await expect( - page.getByRole('menuitem', { name: 'My Data' }) - ).toBeVisible(); - await expect( - page.getByRole('menuitem', { name: 'Following' }) - ).toBeVisible(); - - await page.keyboard.press('Escape'); - await expect(filterMenu).not.toBeVisible(); - }); - } -); + return; + } + + await expect(sortDropdown).toBeVisible(); + + // Click to open dropdown + await sortDropdown.click(); + await page.locator('.ant-dropdown').waitFor({ state: 'visible' }); + + // Verify filter options are present + await expect( + page.getByRole('menuitem', { name: 'All Activity' }) + ).toBeVisible(); + await expect(page.getByRole('menuitem', { name: 'My Data' })).toBeVisible(); + await expect( + page.getByRole('menuitem', { name: 'Following' }) + ).toBeVisible(); + + // Close dropdown by pressing escape + await page.keyboard.press('Escape'); + }); +}); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ActivityFeedTabBadge.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ActivityFeedTabBadge.spec.ts deleted file mode 100644 index ea974d4f2c61..000000000000 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/ActivityFeedTabBadge.spec.ts +++ /dev/null @@ -1,161 +0,0 @@ -/* - * Copyright 2026 Collate. - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -import { expect, test } from '@playwright/test'; -import { TableClass } from '../../support/entity/TableClass'; -import { UserClass } from '../../support/user/UserClass'; -import { performAdminLogin } from '../../utils/admin'; -import { waitForPageLoaded } from '../../utils/polling'; -import { waitForTaskListResponse } from '../../utils/task'; - -async function createOpenTask( - apiContext: Awaited>['apiContext'], - tableFqn: string, - assigneeName: string -): Promise<{ id: string }> { - const res = await apiContext.post('/api/v1/tasks', { - data: { - name: `badge-test-${Date.now()}`, - about: `<#E::table::${tableFqn}>`, - type: 'DescriptionUpdate', - category: 'MetadataUpdate', - assignees: [assigneeName], - }, - }); - expect(res.ok()).toBe(true); - - return res.json(); -} - -async function resolveTask( - apiContext: Awaited>['apiContext'], - taskId: string -) { - const res = await apiContext.post(`/api/v1/tasks/${taskId}/resolve`, { - data: { resolutionType: 'Approved' }, - }); - expect([200, 201]).toContain(res.status()); -} - -async function navigateToTasksPanel(page: import('@playwright/test').Page) { - await page.getByTestId('activity_feed').click(); - await waitForPageLoaded(page); - - const tasksMenuItem = page - .getByTestId('global-setting-left-panel') - .getByText('Tasks'); - - await expect(tasksMenuItem).toBeVisible(); - - await tasksMenuItem.click(); - await waitForPageLoaded(page); -} - -function badge(page: import('@playwright/test').Page) { - return page.getByTestId('left-panel-task-count').getByTestId('filter-count'); -} - -async function switchToClosedFilter(page: import('@playwright/test').Page) { - await page.getByTestId('user-profile-page-task-filter-icon').click(); - const tasksListResponse = waitForTaskListResponse(page); - await page.getByTestId('closed-tasks').click(); - await tasksListResponse; -} - -async function switchToOpenFilter(page: import('@playwright/test').Page) { - await page.getByTestId('user-profile-page-task-filter-icon').click(); - const tasksListResponse = waitForTaskListResponse(page); - await page.getByTestId('open-tasks').click(); - await tasksListResponse; -} -test.describe('ActivityFeedTab — task filter badge and placeholder', () => { - const table = new TableClass(); - const assigneeUser = new UserClass(); - - test.beforeAll(async ({ browser }) => { - const { apiContext, afterAction } = await performAdminLogin(browser); - await table.create(apiContext); - await assigneeUser.create(apiContext); - await afterAction(); - }); - - test.afterAll(async ({ browser }) => { - const { apiContext, afterAction } = await performAdminLogin(browser); - await table.delete(apiContext); - await assigneeUser.delete(apiContext); - await afterAction(); - }); - - test('badge reflects openTaskCount in Open filter and closedTaskCount in Closed filter', async ({ - browser, - }) => { - const { page, apiContext, afterAction } = await performAdminLogin(browser); - - try { - const fqn = table.entityResponseData?.fullyQualifiedName as string; - const task = await createOpenTask( - apiContext, - fqn, - assigneeUser.responseData.name - ); - - await table.visitEntityPage(page); - await navigateToTasksPanel(page); - - await expect(badge(page)).toHaveText('1'); - - await switchToClosedFilter(page); - await expect(badge(page)).toHaveText('0'); - - await resolveTask(apiContext, task.id); - - await page.reload(); - await waitForPageLoaded(page); - await navigateToTasksPanel(page); - - await expect(badge(page)).toHaveText('0'); - - await switchToClosedFilter(page); - await expect(badge(page)).toHaveText('1'); - } finally { - await afterAction(); - } - }); - - test('placeholder shows the correct message per filter state', async ({ - browser, - }) => { - const { page, apiContext, afterAction } = await performAdminLogin(browser); - - const emptyTable = new TableClass(); - - try { - await emptyTable.create(apiContext); - await emptyTable.visitEntityPage(page); - await navigateToTasksPanel(page); - - await expect(page.getByText(/Great News/i)).toBeVisible(); - await switchToClosedFilter(page); - - await expect(page.getByText(/Nothing Closed Yet/i)).toBeVisible(); - await expect(page.getByText(/Great News/i)).not.toBeVisible(); - - await switchToOpenFilter(page); - await expect(page.getByText(/Great News/i)).toBeVisible(); - await expect(page.getByText(/Nothing Closed Yet/i)).not.toBeVisible(); - } finally { - await emptyTable.delete(apiContext); - await afterAction(); - } - }); -}); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/CustomizeDetailPage.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/CustomizeDetailPage.spec.ts index df511c688439..faecba7b3839 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/CustomizeDetailPage.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/CustomizeDetailPage.spec.ts @@ -128,6 +128,8 @@ test.beforeAll('Setup Customize tests', async ({ browser }) => { }); test.afterAll('Cleanup Customize tests', async ({ browser }) => { + test.slow(); + const { apiContext, afterAction } = await performAdminLogin(browser); await adminUser.delete(apiContext); await user.delete(apiContext); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/IncidentManager.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/IncidentManager.spec.ts index 7d27fa7f4cce..a7d58419131c 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/IncidentManager.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/IncidentManager.spec.ts @@ -182,7 +182,7 @@ const openIncidentReassignModal = async (page: Page, testCaseName?: string) => { .last(); const incidentListRowAction = testCaseName ? page - .locator('[data-testid="test-case-incident-manager-table"] tbody tr') + .locator('.ant-table-tbody tr') .filter({ hasText: testCaseName }) .first() .locator('button') diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/KnowledgeCenter.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/KnowledgeCenter.spec.ts index bf3812202af3..10d606f5be3a 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/KnowledgeCenter.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Features/KnowledgeCenter.spec.ts @@ -13,9 +13,8 @@ import { expect, test } from '@playwright/test'; // import cryptoRandomString from 'crypto-random-string-with-promisify-polyfill'; import { SidebarItem } from '../../constant/sidebar'; -import { DataProduct } from '../../support/domain/DataProduct'; -import { Domain } from '../../support/domain/Domain'; import { EntityTypeEndpoint } from '../../support/entity/Entity.interface'; +import { EntityDataClass } from '../../support/entity/EntityDataClass'; import { KnowledgeCenterClass } from '../../support/entity/KnowledgeCenterClass'; import { TableClass } from '../../support/entity/TableClass'; import { TopicClass } from '../../support/entity/TopicClass'; @@ -77,8 +76,6 @@ const knowledgePageQuickLink = { const dataAsset = new TopicClass(); const tableAsset = new TableClass(); const user = new UserClass(); -const domain = new Domain(); -const dataProduct = new DataProduct([domain]); const knowledgeCenter = new KnowledgeCenterClass({}, undefined, tableAsset); const knowledgeCenter1 = new KnowledgeCenterClass(); @@ -95,8 +92,6 @@ test.describe('Knowledge Center', () => { await user.create(apiContext); await tableAsset.create(apiContext); await dataAsset.create(apiContext); - await domain.create(apiContext); - await dataProduct.create(apiContext); await knowledgeCenter.create(apiContext, 15); await knowledgeCenter1.create(apiContext, 2); await afterAction(); @@ -126,16 +121,19 @@ test.describe('Knowledge Center', () => { // add title await addTitle(page, knowledgePageArticle.title); - await assignSingleSelectDomain(page, domain.responseData); + await assignSingleSelectDomain( + page, + EntityDataClass.domain1.responseData + ); await assignDataProduct( page, - domain.responseData, - [dataProduct.responseData], + EntityDataClass.domain1.responseData, + [EntityDataClass.dataProduct1.responseData], 'Add' ); - await removeDataProduct(page, dataProduct.responseData); + await removeDataProduct(page, EntityDataClass.dataProduct1.responseData); // update owner await addMultiOwner({ diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataConsumer.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataConsumer.spec.ts index 35ec6a4af014..0ddb49d214e3 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataConsumer.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataConsumer.spec.ts @@ -203,6 +203,8 @@ entities.forEach((EntityClass) => { }); test.afterAll('Cleanup', async ({ browser }) => { + test.slow(); + const { apiContext, afterAction } = await performAdminLogin(browser); await user.delete(apiContext); await entity.delete(apiContext); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataSteward.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataSteward.spec.ts index 31be94d312e9..98220d150204 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataSteward.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/EntityDataSteward.spec.ts @@ -222,6 +222,8 @@ entities.forEach((EntityClass) => { }); test.afterAll('Cleanup', async ({ browser }) => { + test.slow(); + const { apiContext, afterAction } = await performAdminLogin(browser); await user.delete(apiContext); await entity.delete(apiContext); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/SearchIndexApplication.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/SearchIndexApplication.spec.ts index c7ffae150eab..0f9a72674b62 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/SearchIndexApplication.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/SearchIndexApplication.spec.ts @@ -24,14 +24,6 @@ import { settingClick } from '../../utils/sidebar'; // use the admin user to login test.use({ storageState: 'playwright/.auth/admin.json' }); -const SEARCH_INDEX_APP_NAME = 'SearchIndexingApplication'; -const SUCCESSFUL_RUN_STATUS = /success|completed|activeError/i; - -interface AppRunRecordResponse { - startTime?: number; - status?: string; -} - /** * Installs the Search Indexing Application from the marketplace. * Shared by the "Install application" step and the self-healing guard @@ -169,91 +161,6 @@ const verifyLastExecutionRun = async (page: Page, response: Response) => { } }; -const getLatestRunStartTime = async (page: Page) => { - const { apiContext } = await getApiContext(page); - const response = await apiContext.get( - `/api/v1/apps/name/${SEARCH_INDEX_APP_NAME}/runs/latest` - ); - const run = await getAppRunRecord(response); - - return run?.startTime; -}; - -const getAppRunRecord = async (response: Response) => { - if (!response.ok() || response.status() === 204) { - return undefined; - } - - const body = await response.text(); - - return body ? (JSON.parse(body) as AppRunRecordResponse) : undefined; -}; - -const waitForNewSuccessfulRun = async ( - page: Page, - previousRunStartTime?: number -) => { - const { apiContext } = await getApiContext(page); - let completedRunStartTime: number | undefined; - - await expect - .poll( - async () => { - const response = await apiContext.get( - `/api/v1/apps/name/${SEARCH_INDEX_APP_NAME}/runs/latest` - ); - const run = await getAppRunRecord(response); - - if (run?.startTime === undefined) { - return undefined; - } - - if ( - previousRunStartTime !== undefined && - run.startTime <= previousRunStartTime - ) { - return undefined; - } - - if (run.status && SUCCESSFUL_RUN_STATUS.test(run.status)) { - completedRunStartTime = run.startTime; - } - - return run.status; - }, - { - message: 'Wait for a new successful SearchIndexingApplication run', - intervals: [5_000, 10_000, 15_000, 30_000], - timeout: 300_000, - } - ) - .toEqual(expect.stringMatching(SUCCESSFUL_RUN_STATUS)); - - expect(completedRunStartTime).toBeDefined(); - - return completedRunStartTime; -}; - -const rerunSearchIndexApplicationForTable = async ( - page: Page, - previousRunStartTime?: number -) => { - const { apiContext } = await getApiContext(page); - const response = await apiContext.post( - `/api/v1/apps/trigger/${SEARCH_INDEX_APP_NAME}`, - { - data: { - batchSize: 100, - entities: ['table'], - }, - } - ); - - expect(response.status()).toBeLessThan(300); - - return waitForNewSuccessfulRun(page, previousRunStartTime); -}; - test.describe('Search Index Application', PLAYWRIGHT_BASIC_TEST_TAG_OBJ, () => { test('Search Index Application', async ({ page }) => { test.slow(); @@ -443,14 +350,13 @@ test.describe('Search Index Application', PLAYWRIGHT_BASIC_TEST_TAG_OBJ, () => { }); if (process.env.PLAYWRIGHT_IS_OSS) { - await test.step('Run application and rerun with table-only config', async () => { + await test.step('Run application', async () => { test.slow(true); // Test time shouldn't exceed while re-fetching the history API. await page.click( '[data-testid="search-indexing-application-card"] [data-testid="config-btn"]' ); - const previousRunStartTime = await getLatestRunStartTime(page); const triggerPipelineResponse = page.waitForResponse( '/api/v1/apps/trigger/SearchIndexingApplication' ); @@ -469,12 +375,6 @@ test.describe('Search Index Application', PLAYWRIGHT_BASIC_TEST_TAG_OBJ, () => { expect(statusResponse.status()).toBe(200); await verifyLastExecutionRun(page, statusResponse); - const firstRunStartTime = await waitForNewSuccessfulRun( - page, - previousRunStartTime - ); - - await rerunSearchIndexApplicationForTable(page, firstRunStartTime); }); } }); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/EntityVersionPages.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/EntityVersionPages.spec.ts index e9cee128890b..058fb666eaf3 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/EntityVersionPages.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/EntityVersionPages.spec.ts @@ -71,6 +71,8 @@ const test = base.extend<{ page: Page }>({ test.describe('Entity Version pages', () => { test.beforeAll('Setup pre-requests', async ({ browser }) => { + test.slow(); + adminUser = new UserClass(); entities = entityClasses.map((EntityClass) => new EntityClass()); @@ -142,6 +144,8 @@ test.describe('Entity Version pages', () => { }); test.afterAll('Cleanup', async ({ browser }) => { + test.slow(); + const { apiContext, afterAction } = await performAdminLogin(browser); await adminUser.delete(apiContext); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/ServiceEntityVersionPage.spec.ts b/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/ServiceEntityVersionPage.spec.ts index 4bef42fe841f..15a4e03ab2fb 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/ServiceEntityVersionPage.spec.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/e2e/VersionPages/ServiceEntityVersionPage.spec.ts @@ -64,6 +64,8 @@ const test = base.extend<{ page: Page }>({ test.describe('Service Version pages', () => { test.beforeAll('Setup pre-requests', async ({ browser }) => { + test.slow(); + const { apiContext, afterAction } = await performAdminLogin(browser); await adminUser.create(apiContext); await adminUser.setAdminRole(apiContext); @@ -116,6 +118,8 @@ test.describe('Service Version pages', () => { }); test.afterAll('Cleanup', async ({ browser }) => { + test.slow(); + const { apiContext, afterAction } = await performAdminLogin(browser); await adminUser.delete(apiContext); diff --git a/openmetadata-ui/src/main/resources/ui/playwright/utils/activityAPI.ts b/openmetadata-ui/src/main/resources/ui/playwright/utils/activityAPI.ts deleted file mode 100644 index 62087ad1adf3..000000000000 --- a/openmetadata-ui/src/main/resources/ui/playwright/utils/activityAPI.ts +++ /dev/null @@ -1,336 +0,0 @@ -/* - * Copyright 2025 Collate. - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -import { APIRequestContext, expect, Locator, Page } from '@playwright/test'; -import { TableClass } from '../support/entity/TableClass'; -import { TagClass } from '../support/tag/TagClass'; -import { createAdminApiContext } from './admin'; -import { getApiContext } from './common'; -import { waitForAllLoadersToDisappear } from './entity'; - -export const ACTIVITY_EVENT_TIMEOUT = 300_000; -export const ACTIVITY_TEST_TIMEOUT = ACTIVITY_EVENT_TIMEOUT + 60_000; -export const ACTIVITY_FEED_RESPONSE_TIMEOUT = 15_000; -export const FEED_ITEM_TIMEOUT = 30_000; -export const THUMBS_UP_REACTION = 'thumbsUp'; -export const THUMBS_UP_EMOJI = '👍'; - -const JSON_PATCH_CONTENT_TYPE = 'application/json-patch+json'; - -export type ActivityEventType = - | 'DescriptionUpdated' - | 'OwnerUpdated' - | 'TagsUpdated'; - -export type ActivityApiEvent = Record & { - actor?: { displayName?: string; name?: string }; - eventType?: string; - summary?: string; -}; - -export type ActivityApiResponse = { - data?: ActivityApiEvent[]; -}; - -type FeedThread = { - id?: string; - message?: string; -}; - -type FeedResponse = { - data?: FeedThread[]; -}; - -export const getTableFqn = (table: TableClass) => - table.entityResponseData.fullyQualifiedName ?? ''; - -export const getTableLeafName = (table: TableClass) => - getTableFqn(table).split('.').pop() ?? getTableFqn(table); - -const getTableEntityLink = (table: TableClass) => - `<#E::table::${getTableFqn(table)}>`; - -export const getActivityFeedItems = (page: Page) => - page.locator('#center-container').getByTestId('message-container'); - -export const getFeedItemByText = async (page: Page, text: string) => { - const feedItem = getActivityFeedItems(page).filter({ hasText: text }).first(); - - await expect(feedItem).toBeVisible({ timeout: FEED_ITEM_TIMEOUT }); - await expect(feedItem).toContainText(text); - - return feedItem; -}; - -export const openActivityFeedAndWaitForApi = async ( - page: Page, - entityFqn: string -) => { - const expectedActivityPath = `/api/v1/activity/entity/table/name/${entityFqn}`; - const activityResponsePromise = page.waitForResponse( - (response) => - response.request().method() === 'GET' && - decodeURIComponent(response.url()).includes(expectedActivityPath) && - response.ok(), - { timeout: ACTIVITY_FEED_RESPONSE_TIMEOUT } - ); - - await page.getByTestId('activity_feed').click(); - - const activityResponse = await activityResponsePromise; - - expect(activityResponse.status()).toBe(200); - await waitForAllLoadersToDisappear(page); - - return (await activityResponse.json()) as ActivityApiResponse; -}; - -export const visitTableActivityFeed = async (page: Page, table: TableClass) => { - await table.visitEntityPage(page); - await waitForAllLoadersToDisappear(page); - - return await openActivityFeedAndWaitForApi(page, getTableFqn(table)); -}; - -export const waitForActivityEvent = async ({ - entityFqn, - eventType, - text, -}: { - entityFqn: string; - eventType: ActivityEventType; - text?: string; -}) => { - const { apiContext, afterAction } = await createAdminApiContext(); - const activityUrl = `/api/v1/activity/entity/table/name/${encodeURIComponent( - entityFqn - )}?days=30&limit=50`; - let events: ActivityApiEvent[] = []; - - try { - await expect - .poll( - async () => { - const response = await apiContext.get(activityUrl); - - if (!response.ok()) { - return false; - } - - const body = (await response.json()) as ActivityApiResponse; - events = body.data ?? []; - - return events.some( - (event) => - event.eventType === eventType && - (text === undefined || JSON.stringify(event).includes(text)) - ); - }, - { - timeout: ACTIVITY_EVENT_TIMEOUT, - intervals: [1_000, 2_000, 5_000, 10_000], - message: `Timed out waiting for ${eventType} event for ${entityFqn}`, - } - ) - .toBe(true); - - return events.find( - (event) => - event.eventType === eventType && - (text === undefined || JSON.stringify(event).includes(text)) - ); - } finally { - await afterAction(); - } -}; - -const waitForConversationThread = async ({ - apiContext, - entityLink, - message, - threadId, -}: { - apiContext: APIRequestContext; - entityLink: string; - message: string; - threadId?: string; -}) => { - await expect - .poll( - async () => { - const response = await apiContext.get('/api/v1/feed', { - params: { - entityLink, - type: 'Conversation', - limit: '25', - }, - }); - - if (!response.ok()) { - return false; - } - - const data = (await response.json()) as FeedResponse; - - return (data.data ?? []).some( - (thread) => thread.id === threadId || thread.message === message - ); - }, - { - timeout: 60_000, - intervals: [2_000], - message: `Timed out waiting for conversation "${message}"`, - } - ) - .toBe(true); -}; - -export const createConversationThread = async ( - apiContext: APIRequestContext, - table: TableClass, - message: string -) => { - const entityLink = getTableEntityLink(table); - const response = await apiContext.post('/api/v1/feed', { - data: { - message, - about: entityLink, - }, - }); - - expect(response.ok()).toBeTruthy(); - - const thread = (await response.json()) as FeedThread; - - await waitForConversationThread({ - apiContext, - entityLink, - message, - threadId: thread.id, - }); - - return thread; -}; - -export const createConversationThreadFromPage = async ( - page: Page, - table: TableClass, - message: string -) => { - const { apiContext, afterAction } = await getApiContext(page); - - try { - return await createConversationThread(apiContext, table, message); - } finally { - await afterAction(); - } -}; - -export const patchTableDescription = async ( - apiContext: APIRequestContext, - table: TableClass, - description: string -) => { - const response = await apiContext.patch( - `/api/v1/tables/${table.entityResponseData.id}`, - { - data: [ - { - op: 'add', - path: '/description', - value: description, - }, - ], - headers: { - 'Content-Type': JSON_PATCH_CONTENT_TYPE, - }, - } - ); - - expect(response.ok()).toBeTruthy(); -}; - -export const createDescriptionActivityEventFromPage = async ( - page: Page, - table: TableClass, - description: string -) => { - const { apiContext, afterAction } = await getApiContext(page); - const entityFqn = getTableFqn(table); - - try { - await patchTableDescription(apiContext, table, description); - } finally { - await afterAction(); - } - - return await waitForActivityEvent({ - entityFqn, - eventType: 'DescriptionUpdated', - text: description, - }); -}; - -export const addTagToTable = async ( - apiContext: APIRequestContext, - table: TableClass, - tag: TagClass -) => { - const response = await apiContext.patch( - `/api/v1/tables/${table.entityResponseData.id}`, - { - data: [ - { - op: 'add', - path: '/tags/0', - value: { - tagFQN: tag.responseData.fullyQualifiedName, - source: 'Classification', - }, - }, - ], - headers: { - 'Content-Type': JSON_PATCH_CONTENT_TYPE, - }, - } - ); - - expect(response.ok()).toBeTruthy(); -}; - -export const toggleThumbsUpReaction = async (feedItem: Locator, page: Page) => { - const addReactionButton = feedItem - .getByTestId('feed-reaction-container') - .getByTestId('add-reactions'); - - await expect(addReactionButton).toBeVisible(); - await expect(addReactionButton).toBeEnabled(); - await addReactionButton.click(); - await expect(page.locator('.ant-popover-feed-reactions')).toBeVisible(); - - const reactionResponse = page.waitForResponse( - (response) => - (response.url().includes('/api/v1/activity') || - response.url().includes('/api/v1/feed')) && - response.url().includes(`/reaction/${THUMBS_UP_REACTION}`) && - response.ok() - ); - - await page - .locator(`[data-testid="reaction-button"][title="${THUMBS_UP_REACTION}"]`) - .click(); - - const response = await reactionResponse; - - expect(response.ok()).toBeTruthy(); - await waitForAllLoadersToDisappear(page); -}; diff --git a/openmetadata-ui/src/main/resources/ui/playwright/utils/incidentManager.ts b/openmetadata-ui/src/main/resources/ui/playwright/utils/incidentManager.ts index 3a7f344c21d3..7df393c378e3 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/utils/incidentManager.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/utils/incidentManager.ts @@ -74,13 +74,15 @@ export const addAssigneeFromPopoverWidget = async (data: { if (testCaseName) { const incidentRow = page - .locator('tr') - .filter({ has: page.getByTestId(`test-case-${testCaseName}`) }) + .getByRole('row', { name: new RegExp(testCaseName, 'i') }) .first(); const editOwnerButton = incidentRow.getByTestId('edit-owner'); - await expect(editOwnerButton).toBeVisible(); - await editOwnerButton.click(); + if (await editOwnerButton.isVisible().catch(() => false)) { + await editOwnerButton.click(); + } else { + await incidentRow.locator('td').last().getByRole('button').click(); + } } else if (await taskTabEditAssigneesButton.isVisible().catch(() => false)) { await taskTabEditAssigneesButton.click(); await waitForAllLoadersToDisappear(page); @@ -182,7 +184,7 @@ export const assignIncident = async (data: { .poll( async () => { const incidentRow = page - .getByTestId(`test-case-${testCaseName}`) + .getByRole('row', { name: new RegExp(testCaseName, 'i') }) .first(); const incidentLink = page .getByRole('link', { name: testCaseName }) diff --git a/openmetadata-ui/src/main/resources/ui/playwright/utils/metric.ts b/openmetadata-ui/src/main/resources/ui/playwright/utils/metric.ts index 86dc7f734266..a42e517ab6f8 100644 --- a/openmetadata-ui/src/main/resources/ui/playwright/utils/metric.ts +++ b/openmetadata-ui/src/main/resources/ui/playwright/utils/metric.ts @@ -13,7 +13,7 @@ import { expect, Page } from '@playwright/test'; import { EntityTypeEndpoint } from '../support/entity/Entity.interface'; import { MetricClass } from '../support/entity/MetricClass'; -import { clickOutside, descriptionBox, uuid } from './common'; +import { descriptionBox, uuid } from './common'; import { hardDeleteEntity, waitForAllLoadersToDisappear } from './entity'; export const updateMetricType = async (page: Page, metric: string) => { @@ -175,8 +175,6 @@ export const updateRelatedMetric = async ( }) .click(); - // perform click outside to close the select options and make click to button - await clickOutside(page); await page.locator('[data-testid="saveRelatedMetrics"]').click(); await patchPromise; diff --git a/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md b/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md index e56d1d4058e0..7f30ac08872f 100644 --- a/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md +++ b/openmetadata-ui/src/main/resources/ui/public/locales/en-US/Applications/SearchIndexingApplication.md @@ -70,10 +70,15 @@ $$section $$ +$$section +### Recreate Indexes $(id="recreateIndex") + +$$ + $$section ### Search Index Language $(id="searchIndexMappingLanguage") -Search index mapping language. +Recreate Indexes with updated Language $$ @@ -84,6 +89,13 @@ Enable automatic performance tuning based on cluster capabilities and database e $$ +$$section +### Use Distributed Indexing $(id="useDistributedIndexing") + +Enable distributed indexing to scale reindexing across multiple servers with fault tolerance and parallel processing + +$$ + $$section ### Partition Size $(id="partitionSize") diff --git a/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.component.test.tsx b/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.component.test.tsx deleted file mode 100644 index 676725d3d700..000000000000 --- a/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.component.test.tsx +++ /dev/null @@ -1,283 +0,0 @@ -/* - * Copyright 2026 Collate. - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * http://www.apache.org/licenses/LICENSE-2.0 - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -import { fireEvent, render, screen, waitFor } from '@testing-library/react'; -import { MemoryRouter } from 'react-router-dom'; -import { EntityType } from '../../../enums/entity.enum'; -import { FeedFilter } from '../../../enums/mydata.enum'; -import { ActivityFeedTab } from './ActivityFeedTab.component'; -import { - ActivityFeedLayoutType, - ActivityFeedTabs, -} from './ActivityFeedTab.interface'; - -const mockGetFeedData = jest.fn(); -const mockGetTaskData = jest.fn(); -const mockGetTaskCounts = jest.fn(); -const mockUseRequiredParams = jest.fn(); - -jest.mock('../../../hooks/useApplicationStore', () => ({ - useApplicationStore: () => ({ - currentUser: { id: 'u1', name: 'admin', fullyQualifiedName: 'admin' }, - }), -})); - -jest.mock('../../../hooks/authHooks', () => ({ - useAuth: () => ({ isAdminUser: false }), -})); - -jest.mock('../../../hooks/useDomainStore', () => ({ - useDomainStore: (selector: (s: { activeDomain: string }) => string) => - selector({ activeDomain: 'All Domains' }), -})); - -jest.mock('../../../hooks/useFqn', () => ({ - useFqn: () => ({ fqn: 'test.db.table' }), -})); - -jest.mock('../../../utils/useRequiredParams', () => ({ - useRequiredParams: () => mockUseRequiredParams(), -})); - -jest.mock('../../../hooks/useElementInView', () => ({ - useElementInView: () => [{ current: null }, false], -})); - -jest.mock('../ActivityFeedProvider/ActivityFeedProvider', () => ({ - useActivityFeedProvider: () => ({ - selectedThread: null, - setActiveThread: jest.fn(), - entityThread: [], - getFeedData: mockGetFeedData, - getTaskData: mockGetTaskData, - loading: false, - entityPaging: {}, - tasks: [], - selectedTask: null, - setActiveTask: jest.fn(), - activityEvents: [], - isActivityLoading: false, - fetchEntityActivity: jest.fn(), - fetchUserActivity: jest.fn(), - userId: '', - selectedActivity: null, - setActiveActivity: jest.fn(), - }), -})); - -jest.mock('../../../rest/tasksAPI', () => ({ - getTaskCounts: (...args: unknown[]) => mockGetTaskCounts(...args), - TaskStatusGroup: { Open: 'open', Closed: 'closed' }, -})); - -jest.mock('../../../rest/feedsAPI', () => ({ - getFeedCount: jest - .fn() - .mockResolvedValue([{ conversationCount: 0, mentionCount: 0 }]), -})); - -jest.mock('../../../utils/CommonUtils', () => ({ - getCountBadge: (count: number) => ( - {count} - ), - getFeedCounts: jest.fn((_, __, ___, cb) => - cb({ - conversationCount: 0, - mentionCount: 0, - totalCount: 0, - totalTasksCount: 0, - openTaskCount: 0, - closedTaskCount: 0, - }) - ), -})); - -jest.mock('../../../utils/ToastUtils', () => ({ - showErrorToast: jest.fn(), -})); - -jest.mock('../../../utils/EntityUtilClassBase', () => ({ - default: { getActivityFeedTabs: jest.fn().mockReturnValue([]) }, -})); - -jest.mock('../../../utils/EntityUtils', () => ({ - getEntityUserLink: jest.fn().mockReturnValue(''), -})); - -jest.mock('../ActivityFeedList/ActivityFeedListV1New.component', () => - jest.fn().mockReturnValue(
) -); - -jest.mock('../ActivityFeedList/TaskListV1.component', () => - jest - .fn() - .mockImplementation(({ emptyPlaceholderText }) => ( -
{emptyPlaceholderText}
- )) -); - -jest.mock('../ActivityFeedPanel/FeedPanelBodyV1New', () => - jest.fn().mockReturnValue(
) -); - -jest.mock('../../Entity/Task/TaskTab/TaskTabNew.component', () => ({ - TaskTabNew: jest.fn().mockReturnValue(
), -})); - -jest.mock('../../common/ErrorWithPlaceholder/ErrorPlaceHolderNew', () => - jest.fn().mockReturnValue(
) -); - -jest.mock('../../common/Loader/Loader', () => - jest.fn().mockReturnValue(
) -); - -jest.mock('../../MyData/Widgets/FeedsWidget/feeds-widget.less', () => ({})); -jest.mock('./activity-feed-tab.less', () => ({})); - -const defaultProps = { - entityType: EntityType.TABLE as EntityType.TABLE, - onFeedUpdate: jest.fn(), - layoutType: ActivityFeedLayoutType.THREE_PANEL, -}; - -const renderComponent = (subTab = ActivityFeedTabs.TASKS) => { - mockUseRequiredParams.mockReturnValue({ tab: 'activity_feed', subTab }); - - return render( - - - - ); -}; - -describe('ActivityFeedTab', () => { - beforeEach(() => { - jest.clearAllMocks(); - mockGetTaskCounts.mockResolvedValue({ - open: 0, - inProgress: 0, - completed: 0, - total: 0, - }); - mockGetFeedData.mockResolvedValue(undefined); - mockGetTaskData.mockResolvedValue(undefined); - }); - - describe('Bug 1 — feedFilter uses ActivityFeedTabs.MENTIONS enum', () => { - it('calls getFeedData with FeedFilter.MENTIONS when mentions tab is active', async () => { - renderComponent(ActivityFeedTabs.MENTIONS); - - await waitFor(() => { - const calls = mockGetFeedData.mock.calls; - const mentionsCall = calls.find( - ([feedFilter]) => feedFilter === FeedFilter.MENTIONS - ); - - expect(mentionsCall).toBeDefined(); - }); - }); - - it('does not call getFeedData with FeedFilter.MENTIONS when tasks tab is active', async () => { - renderComponent(ActivityFeedTabs.TASKS); - - await waitFor(() => expect(mockGetTaskData).toHaveBeenCalled()); - - const mentionsCall = mockGetFeedData.mock.calls.find( - ([feedFilter]) => feedFilter === FeedFilter.MENTIONS - ); - - expect(mentionsCall).toBeUndefined(); - }); - }); - - describe('Bug 3/4 — badge and placeholder reflect taskFilter state', () => { - it('left-panel badge shows openTaskCount in Open filter', async () => { - mockGetTaskCounts.mockResolvedValue({ - open: 3, - inProgress: 0, - completed: 5, - total: 8, - }); - - renderComponent(ActivityFeedTabs.TASKS); - - await waitFor(() => { - const badge = screen.getByTestId('left-panel-task-count'); - - expect(badge).toHaveTextContent('3'); - }); - }); - - it('left-panel badge switches to closedTaskCount when Closed filter is selected', async () => { - mockGetTaskCounts.mockResolvedValue({ - open: 3, - inProgress: 0, - completed: 5, - total: 8, - }); - - renderComponent(ActivityFeedTabs.TASKS); - - await waitFor(() => - expect(screen.getByTestId('left-panel-task-count')).toHaveTextContent( - '3' - ) - ); - - fireEvent.click(screen.getByTestId('user-profile-page-task-filter-icon')); - - const closedItem = await screen.findByTestId('closed-tasks'); - - fireEvent.click(closedItem); - - await waitFor(() => - expect(screen.getByTestId('left-panel-task-count')).toHaveTextContent( - '5' - ) - ); - }); - - it('placeholder shows open tasks message when Open filter is active', async () => { - renderComponent(ActivityFeedTabs.TASKS); - - await waitFor(() => { - expect( - screen.getByText('message.no-open-tasks-title') - ).toBeInTheDocument(); - }); - }); - - it('placeholder shows closed tasks message when Closed filter is selected', async () => { - renderComponent(ActivityFeedTabs.TASKS); - - await waitFor(() => - expect( - screen.getByTestId('user-profile-page-task-filter-icon') - ).toBeInTheDocument() - ); - - fireEvent.click(screen.getByTestId('user-profile-page-task-filter-icon')); - - const closedItem = await screen.findByTestId('closed-tasks'); - - fireEvent.click(closedItem); - - await waitFor(() => { - expect( - screen.getByText('message.no-closed-tasks-title') - ).toBeInTheDocument(); - }); - }); - }); -}); diff --git a/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.component.tsx index 52edcaea0e79..0d04351162d4 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.component.tsx +++ b/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.component.tsx @@ -118,9 +118,7 @@ export const ActivityFeedTab = ({ tab: EntityTabs; subTab: ActivityFeedTabs; }>(); - const [taskFilter, setTaskFilter] = useState( - TaskStatusGroup.Open - ); + const [taskFilter, setTaskFilter] = useState('open'); const [isFullWidth, setIsFullWidth] = useState(false); const [countData, setCountData] = useState<{ loading: boolean; @@ -211,17 +209,6 @@ export const ActivityFeedTab = ({ {t('message.no-mentions')} ); - } else if (taskFilter === TaskStatusGroup.Closed) { - return ( -
- - {t('message.no-closed-tasks-title')} - - - {t('message.no-closed-tasks-description')} - -
- ); } else { return (
@@ -234,7 +221,7 @@ export const ActivityFeedTab = ({
); } - }, [activeTab, taskFilter, t]); + }, [activeTab]); const handleFeedCount = useCallback( (data: FeedCounts) => { @@ -307,8 +294,7 @@ export const ActivityFeedTab = ({ activeTab === ActivityFeedTabs.ALL ? ThreadType.Conversation : undefined, - feedFilter: - activeTab === ActivityFeedTabs.MENTIONS ? FeedFilter.MENTIONS : filter, + feedFilter: activeTab === 'mentions' ? FeedFilter.MENTIONS : filter, }; }, [activeTab, isAdminUser, currentUser, fqn, isUserEntity]); @@ -488,16 +474,16 @@ export const ActivityFeedTab = ({ const taskFilterOptions = useMemo( () => [ { - key: TaskStatusGroup.Open, + key: 'open', label: (
- {taskFilter === TaskStatusGroup.Open ? ( + {taskFilter === 'open' ? ( {t('label.open')}
{countData?.data?.openTaskCount} @@ -523,21 +509,21 @@ export const ActivityFeedTab = ({
), onClick: () => { - handleUpdateTaskFilter(TaskStatusGroup.Open); + handleUpdateTaskFilter('open'); setActiveTask(); }, }, { - key: TaskStatusGroup.Closed, + key: 'closed', label: (
- {taskFilter === TaskStatusGroup.Closed ? ( + {taskFilter === 'closed' ? ( {t('label.closed')}
{countData?.data?.closedTaskCount} @@ -566,7 +552,7 @@ export const ActivityFeedTab = ({
), onClick: () => { - handleUpdateTaskFilter(TaskStatusGroup.Closed); + handleUpdateTaskFilter('closed'); setActiveTask(); }, }, @@ -578,6 +564,7 @@ export const ActivityFeedTab = ({ return ( handleTabChange(value as ActivityFeedTabs)} /> ); @@ -753,9 +739,7 @@ export const ActivityFeedTab = ({ {getCountBadge( - taskFilter === TaskStatusGroup.Open - ? countData?.data?.openTaskCount - : countData?.data?.closedTaskCount, + countData?.data?.openTaskCount, '', isTaskActiveTab )} @@ -803,7 +787,7 @@ export const ActivityFeedTab = ({ - {taskFilter === TaskStatusGroup.Open + {taskFilter === 'open' ? `${t('label.open')} (${ countData?.data?.openTaskCount ?? 0 })` diff --git a/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.interface.ts b/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.interface.ts index e8894eefb791..b590a9c8cdd1 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.interface.ts +++ b/openmetadata-ui/src/main/resources/ui/src/components/ActivityFeed/ActivityFeedTab/ActivityFeedTab.interface.ts @@ -15,6 +15,8 @@ import { Column } from '../../../generated/entity/data/table'; import { EntityReference } from '../../../generated/entity/type'; import { FeedCounts } from '../../../interface/feed.interface'; +export type TaskFilter = 'open' | 'close'; + export enum ActivityFeedTabs { ALL = 'all', MENTIONS = 'mentions', diff --git a/openmetadata-ui/src/main/resources/ui/src/components/AppRouter/AuthenticatedAppRouter.tsx b/openmetadata-ui/src/main/resources/ui/src/components/AppRouter/AuthenticatedAppRouter.tsx index a04894127331..bca763ae4c30 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/AppRouter/AuthenticatedAppRouter.tsx +++ b/openmetadata-ui/src/main/resources/ui/src/components/AppRouter/AuthenticatedAppRouter.tsx @@ -199,6 +199,12 @@ const OntologyExplorerPage = withSuspenseFallback( ) ); +const SparqlPlaygroundPage = withSuspenseFallback( + React.lazy( + () => import('../../pages/SparqlPlayground/SparqlPlayground.component') + ) +); + const WorkflowsListPage = withSuspenseFallback( React.lazy( () => @@ -394,6 +400,10 @@ const AuthenticatedAppRouter: FunctionComponent = () => { element={} path={ROUTES.EXPLORE_WITH_TAB} /> + } + path={ROUTES.SPARQL_PLAYGROUND} + /> } path={ROUTES.ONTOLOGY_EXPLORER} diff --git a/openmetadata-ui/src/main/resources/ui/src/components/BlockEditor/Extensions/text-highlight-view.ts b/openmetadata-ui/src/main/resources/ui/src/components/BlockEditor/Extensions/text-highlight-view.ts index a38ea7410774..e66503922771 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/BlockEditor/Extensions/text-highlight-view.ts +++ b/openmetadata-ui/src/main/resources/ui/src/components/BlockEditor/Extensions/text-highlight-view.ts @@ -22,20 +22,10 @@ export default Node.create({ return { class: { default: '', - parseHTML: (element) => element.getAttribute('class'), - renderHTML: (attributes) => { - if (!attributes.class) { - return {}; - } - - return { - class: attributes.class, - }; - }, }, 'data-testid': { default: '', - parseHTML: (element) => element.dataset.testid, + parseHTML: (element) => element.getAttribute('data-testid'), renderHTML: (attributes) => { if (!attributes['data-testid']) { return {}; @@ -48,7 +38,7 @@ export default Node.create({ }, 'data-highlight': { default: true, - parseHTML: (element) => element.dataset.highlight, + parseHTML: (element) => element.getAttribute('data-highlight'), renderHTML: (attributes) => { if (!attributes['data-highlight']) { return {}; @@ -67,9 +57,6 @@ export default Node.create({ { tag: 'span[data-highlight]', }, - { - tag: 'span.text-highlighter', - }, ]; }, @@ -80,7 +67,7 @@ export default Node.create({ textHighlightNode.setAttribute(key, HTMLAttributes[key]); }); - textHighlightNode.dataset.highlight = 'true'; + textHighlightNode.setAttribute('data-highlight', 'true'); textHighlightNode.innerHTML = node.textContent; return { diff --git a/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.component.tsx index df4f6805ff4e..1a4d25f89454 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.component.tsx +++ b/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.component.tsx @@ -65,7 +65,6 @@ import { triggerOnDemandApp } from '../../../rest/applicationAPI'; import { getContractByEntityId } from '../../../rest/contractAPI'; import { getDataQualityLineage } from '../../../rest/lineageAPI'; import { getContainerAncestors } from '../../../rest/storageAPI'; -import { hasEditAccess } from '../../../utils/CommonUtils'; import { getDataAssetsHeaderInfo, isDataAssetsWithServiceField, @@ -581,13 +580,8 @@ export const DataAssetsHeader = ({ ]); const isOwner = useMemo( - () => - Boolean( - currentUser && - dataAsset.owners?.length && - hasEditAccess(dataAsset.owners, currentUser) - ), - [dataAsset.owners, currentUser] + () => dataAsset.owners?.some((o) => o.id === USER_ID) ?? false, + [dataAsset.owners, USER_ID] ); const requestDataAccessButton = useMemo(() => { diff --git a/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.test.tsx b/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.test.tsx index d0e2c6c2cbe3..66995121509a 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.test.tsx +++ b/openmetadata-ui/src/main/resources/ui/src/components/DataAssets/DataAssetsHeader/DataAssetsHeader.test.tsx @@ -747,37 +747,6 @@ describe('DataAssetsHeader component', () => { ).not.toBeInTheDocument(); }); - it('should not render when user belongs to an owner team', async () => { - const { useApplicationStore } = jest.requireMock( - '../../../hooks/useApplicationStore' - ); - (useApplicationStore as jest.Mock).mockReturnValue({ - currentUser: { - id: 'user-2', - name: 'team.member', - teams: [{ id: 'team-1', type: 'team' }], - }, - }); - - render( - - ); - - expect( - screen.queryByTestId('request-data-access-button') - ).not.toBeInTheDocument(); - - (useApplicationStore as jest.Mock).mockReturnValue({ - currentUser: { id: 'user-1', name: 'test.user' }, - }); - }); - it('should render enabled button when no existing DAR task', async () => { mockListTasks.mockResolvedValue({ data: [] }); @@ -790,6 +759,22 @@ describe('DataAssetsHeader component', () => { }); }); + it('should call listTasks with correct server-side filters', async () => { + render(); + + await waitFor(() => { + expect(mockListTasks).toHaveBeenCalledWith( + expect.objectContaining({ + aboutEntity: 'service.db.schema.my_table', + category: 'DataAccess', + type: 'DataAccessRequest', + createdBy: 'test.user', + statusGroup: 'open', + }) + ); + }); + }); + it('should disable button when a task is in review stage', async () => { mockListTasks.mockResolvedValue({ data: [ diff --git a/openmetadata-ui/src/main/resources/ui/src/components/DataProducts/DataProductsDetailsPage/DataProductsDetailsPage.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/DataProducts/DataProductsDetailsPage/DataProductsDetailsPage.component.tsx index 52ab46f02b9e..cba73af3e9a6 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/DataProducts/DataProductsDetailsPage/DataProductsDetailsPage.component.tsx +++ b/openmetadata-ui/src/main/resources/ui/src/components/DataProducts/DataProductsDetailsPage/DataProductsDetailsPage.component.tsx @@ -74,7 +74,6 @@ import { searchQuery } from '../../../rest/searchAPI'; import { getEntityDeleteMessage, getFeedCounts, - hasEditAccess, } from '../../../utils/CommonUtils'; import { checkIfExpandViewSupported, @@ -350,13 +349,8 @@ const DataProductsDetailsPage = ({ ); const isOwner = useMemo( - () => - Boolean( - currentUser && - dataProduct.owners?.length && - hasEditAccess(dataProduct.owners, currentUser) - ), - [dataProduct.owners, currentUser] + () => dataProduct.owners?.some((o) => o.id === currentUser?.id) ?? false, + [dataProduct.owners, currentUser?.id] ); const handleVoteChange = useCallback( diff --git a/openmetadata-ui/src/main/resources/ui/src/components/DataQuality/IncidentManager/DimensionalityTab/DimensionalityTab.tsx b/openmetadata-ui/src/main/resources/ui/src/components/DataQuality/IncidentManager/DimensionalityTab/DimensionalityTab.tsx index 31260c8b3ea7..93d4d9889d24 100644 --- a/openmetadata-ui/src/main/resources/ui/src/components/DataQuality/IncidentManager/DimensionalityTab/DimensionalityTab.tsx +++ b/openmetadata-ui/src/main/resources/ui/src/components/DataQuality/IncidentManager/DimensionalityTab/DimensionalityTab.tsx @@ -10,11 +10,13 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -import { Select, Skeleton, Table } from '@openmetadata/ui-core-components'; +import { Select, Skeleton, Tooltip } from '@openmetadata/ui-core-components'; +import { HelpCircle } from '@untitledui/icons'; +import { ColumnsType } from 'antd/lib/table'; import { format } from 'date-fns'; import { isEmpty, split, toLower } from 'lodash'; import { DateRangeObject } from 'Models'; -import { useCallback, useEffect, useMemo, useState } from 'react'; +import { useEffect, useMemo, useState } from 'react'; import { Trans, useTranslation } from 'react-i18next'; import { Link } from 'react-router-dom'; import { @@ -41,6 +43,7 @@ import NoDataPlaceholderNew from '../../../common/ErrorWithPlaceholder/NoDataPla import MuiDatePickerMenu from '../../../common/MuiDatePickerMenu/MuiDatePickerMenu'; import StatusBadge from '../../../common/StatusBadge/StatusBadge.component'; import { StatusType } from '../../../common/StatusBadge/StatusBadge.interface'; +import Table from '../../../common/Table/Table'; import { ProfilerTabPath } from '../../../Database/Profiler/ProfilerDashboard/profilerDashboard.interface'; import DimensionalityHeatmap from './DimensionalityHeatmap/DimensionalityHeatmap.component'; import { DimensionResultWithTimestamp } from './DimensionalityHeatmap/DimensionalityHeatmap.interface'; @@ -173,65 +176,84 @@ const DimensionalityTab = () => { }); }, [dimensionData]); - const dimensionTableColumns = useMemo( - () => [ - { id: 'status', label: t('label.status') }, - { - id: 'impactScore', - label: t('label.impact-score'), - tooltip: t('message.impact-score-helper'), + const tableColumns: ColumnsType<{ + key: string; + dimensionValue: string; + result: DimensionResultWithTimestamp; + }> = [ + { + title: t('label.status'), + dataIndex: 'result', + key: 'status', + width: 120, + render: (result: DimensionResultWithTimestamp) => { + return result?.testCaseStatus ? ( + + ) : ( + -- + ); }, - { id: 'dimensionValue', label: t('label.dimension') }, - { id: 'lastRun', label: t('label.last-run') }, - ], - [t] - ); - - const renderCell = useCallback( - ( - col: { id: string }, - row: (typeof getLatestResultPerDimension)[number] - ) => { - switch (col.id) { - case 'status': - return row.result?.testCaseStatus ? ( - - ) : ( - -- - ); - case 'impactScore': - return ( - - {row.result?.impactScore ?? '--'} - - ); - case 'dimensionValue': - return ( - - {row.dimensionValue} - - ); - case 'lastRun': - return row.result?.timestamp ? ( - - ) : ( - -- - ); - default: - return null; - } }, - [testCase?.fullyQualifiedName] - ); + { + title: ( +
+ {t('label.impact-score')} + + {t('message.impact-score-helper')} + + }> + + +
+ ), + dataIndex: 'result', + key: 'impactScore', + width: 120, + render: (result: DimensionResultWithTimestamp) => { + return ( + {result?.impactScore ?? '--'} + ); + }, + }, + { + title: t('label.dimension'), + dataIndex: 'dimensionValue', + key: 'dimensionValue', + width: 200, + render: (dimensionValue: string, record) => { + return ( + + {dimensionValue} + + ); + }, + }, + { + title: t('label.last-run'), + dataIndex: 'result', + key: 'lastRun', + width: 200, + render: (result: DimensionResultWithTimestamp) => { + return result?.timestamp ? ( + + ) : ( + -- + ); + }, + }, + ]; const noDataPlaceholder = useMemo(() => { if (isLoading) { @@ -265,9 +287,8 @@ const DimensionalityTab = () => { {`${t('label.select-dimension')}:`}

({ + id: o.value, + label: o.label, + }))} + size="sm" + value={format} + onChange={(key) => + setFormat(String(key) as SparqlPlaygroundFormat) + } + /> +