From 8d7db583049d8f93e0c1330255a2c43588f6eaff Mon Sep 17 00:00:00 2001 From: Ram Narayan Balaji <81347100+yan-3005@users.noreply.github.com> Date: Mon, 4 Aug 2025 18:11:39 +0530 Subject: [PATCH] Fix #22623 Fetch entities for import csv should exclude only the fields that are coming from the CSV (#22663) * Fetch entities for import csv should exclude only the fields that are coming from the CSV * Updated DocStrings with Comments! --- .../java/org/openmetadata/csv/EntityCsv.java | 18 ++- .../java/org/openmetadata/service/Entity.java | 18 +++ .../service/jdbi3/EntityRepository.java | 13 ++ .../openmetadata/service/util/EntityUtil.java | 20 +++ .../databases/DatabaseSchemaResourceTest.java | 141 ++++++++++++++++++ 5 files changed, 204 insertions(+), 6 deletions(-) diff --git a/openmetadata-service/src/main/java/org/openmetadata/csv/EntityCsv.java b/openmetadata-service/src/main/java/org/openmetadata/csv/EntityCsv.java index 4f75d266b90..5a3b63b7b41 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/csv/EntityCsv.java +++ b/openmetadata-service/src/main/java/org/openmetadata/csv/EntityCsv.java @@ -1069,8 +1069,11 @@ public abstract class EntityCsv { Database database; try { database = - Entity.getEntityByName( - DATABASE, dbFQN, "name,displayName,fullyQualifiedName,service", Include.NON_DELETED); + Entity.getEntityByNameWithExcludedFields( + DATABASE, + dbFQN, + "name,displayName,description,owners,tags,glossaryTerms,tiers,certification,retentionPeriod,sourceUrl,domains,extension,updatedAt,updatedBy", + Include.NON_DELETED); } catch (EntityNotFoundException ex) { LOG.warn("Database not found: {}. Handling based on dryRun mode.", dbFQN); if (importResult.getDryRun()) { @@ -1087,10 +1090,10 @@ public abstract class EntityCsv { String schemaFqn = FullyQualifiedName.add(dbFQN, csvRecord.get(0)); try { schema = - Entity.getEntityByName( + Entity.getEntityByNameWithExcludedFields( DATABASE_SCHEMA, schemaFqn, - "name,displayName,fullyQualifiedName", + "name,displayName,description,owners,tags,glossaryTerms,tiers,certification,retentionPeriod,sourceUrl,domains,extension,updatedAt,updatedBy", Include.NON_DELETED); } catch (Exception ex) { LOG.warn("Database Schema not found: {}, it will be created with Import.", schemaFqn); @@ -1169,8 +1172,11 @@ public abstract class EntityCsv { try { table = - Entity.getEntityByName( - TABLE, tableFqn, "name,displayName,fullyQualifiedName,columns", Include.NON_DELETED); + Entity.getEntityByNameWithExcludedFields( + TABLE, + tableFqn, + "name,displayName,description,owners,tags,glossaryTerms,tiers,certification,retentionPeriod,sourceUrl,domains,extension,updatedAt,updatedBy", + Include.NON_DELETED); } catch (EntityNotFoundException ex) { // Table not found, create a new one diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java b/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java index 9985e529796..2c5fbec46a7 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/Entity.java @@ -533,6 +533,24 @@ public final class Entity { return getEntityByName(entityType, fqn, fields, include, true); } + public static T getEntityByNameWithExcludedFields( + String entityType, String fqn, String excludeFields, Include include) { + return getEntityByNameWithExcludedFields(entityType, fqn, excludeFields, include, false); + } + + // Retrieve the entity by name excluding specific fields. Useful for import using CSV where + // certain fields are already sent in the csv + public static T getEntityByNameWithExcludedFields( + String entityType, String fqn, String excludeFields, Include include, boolean fromCache) { + EntityRepository entityRepository = Entity.getEntityRepository(entityType); + @SuppressWarnings("unchecked") + T entity = + (T) + entityRepository.getByNameWithExcludedFields( + null, fqn, excludeFields, include, fromCache); + return entity; + } + public static List getEntityByNames( String entityType, List tagFQNs, String fields, Include include) { EntityRepository entityRepository = Entity.getEntityRepository(entityType); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java index bb1509c69a4..138e41faa99 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java @@ -787,6 +787,19 @@ public abstract class EntityRepository { return entities; } + // Get Entity By Name excluding certain fields + public final T getByNameWithExcludedFields( + UriInfo uriInfo, String fqn, String excludeFields, Include include) { + return getByNameWithExcludedFields(uriInfo, fqn, excludeFields, include, false); + } + + // Form the Field Object by excluding certain fields and get entity by name + public final T getByNameWithExcludedFields( + UriInfo uriInfo, String fqn, String excludeFields, Include include, boolean fromCache) { + Fields fields = EntityUtil.Fields.createWithExcludedFields(allowedFields, excludeFields); + return getByName(uriInfo, fqn, fields, include, fromCache); + } + public final T findByNameOrNull(String fqn, Include include) { try { return findByName(fqn, include); diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/EntityUtil.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/EntityUtil.java index a8de8ea208e..4f61c94d8d1 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/EntityUtil.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/EntityUtil.java @@ -388,6 +388,26 @@ public final class EntityUtil { fieldList.add(field); } + // Create Fields Objects by excluding certain fields + public static Fields createWithExcludedFields( + Set allowedFields, Set excludeFields) { + Set resultFields = new HashSet<>(allowedFields); + if (excludeFields != null) { + resultFields.removeAll(excludeFields); + } + return new Fields(allowedFields, resultFields); + } + + public static Fields createWithExcludedFields( + Set allowedFields, String excludeFieldsParam) { + Set excludeFields = new HashSet<>(); + if (!nullOrEmpty(excludeFieldsParam)) { + excludeFields = + new HashSet<>(Arrays.asList(excludeFieldsParam.replace(" ", "").split(","))); + } + return createWithExcludedFields(allowedFields, excludeFields); + } + @Override public String toString() { return String.join(",", fieldList); diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseSchemaResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseSchemaResourceTest.java index 2bc2567f22e..55bf45db83c 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseSchemaResourceTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseSchemaResourceTest.java @@ -15,6 +15,7 @@ package org.openmetadata.service.resources.databases; import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST; import static org.apache.commons.lang.StringEscapeUtils.escapeCsv; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; @@ -38,6 +39,7 @@ import jakarta.ws.rs.client.WebTarget; import jakarta.ws.rs.core.Response; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -72,6 +74,7 @@ import org.openmetadata.schema.type.EntityReference; import org.openmetadata.schema.type.TableConstraint; import org.openmetadata.schema.type.TableProfilerConfig; import org.openmetadata.schema.type.csv.CsvImportResult; +import org.openmetadata.schema.utils.JsonUtils; import org.openmetadata.service.Entity; import org.openmetadata.service.jdbi3.DatabaseSchemaRepository; import org.openmetadata.service.resources.EntityResourceTest; @@ -685,6 +688,144 @@ public class DatabaseSchemaResourceTest "At least one column should have updated description"); } + @Test + void testImportExportWithTableConstraints() throws IOException { + // Create a schema for this test to avoid conflicts + CreateDatabaseSchema createSchema = createRequest("constraint_test_schema"); + DatabaseSchema schema = createEntity(createSchema, ADMIN_AUTH_HEADERS); + + TableResourceTest tableTest = new TableResourceTest(); + + // Create tables and columns for FK relationships + Column c1 = new Column().withName("user_ref").withDataType(ColumnDataType.STRING); + Column c2 = new Column().withName("tenant_id").withDataType(ColumnDataType.STRING); + Column c3 = new Column().withName("user_id").withDataType(ColumnDataType.STRING); + + // Create target table (referenced table with 2 columns) + Table targetTable = + tableTest.createEntity( + tableTest + .createRequest("target_table") + .withDatabaseSchema(schema.getFullyQualifiedName()) + .withTableConstraints(null) + .withColumns(List.of(c2, c3)), + ADMIN_AUTH_HEADERS); + + // Create source table (no constraints initially) + Table sourceTable = + tableTest.createEntity( + tableTest + .createRequest("source_table") + .withDatabaseSchema(schema.getFullyQualifiedName()) + .withColumns(List.of(c1)) + .withTableConstraints(null), + ADMIN_AUTH_HEADERS); + + // Resolve column FQNs needed for FK definitions + Table targetRef = + tableTest.getEntityByName(targetTable.getFullyQualifiedName(), ADMIN_AUTH_HEADERS); + + // Create foreign key constraint - simple 1:1 mapping + String targetCol1FQN = targetRef.getColumns().getFirst().getFullyQualifiedName(); + + String originalJson = JsonUtils.pojoToJson(sourceTable); + Table sourceTableV2 = JsonUtils.deepCopy(sourceTable, Table.class); + + // Create a simple 1:1 foreign key constraint: 1 local column referencing 1 referred column + TableConstraint foreignKeyConstraint = + new TableConstraint() + .withConstraintType(TableConstraint.ConstraintType.FOREIGN_KEY) + .withColumns(List.of("user_ref")) // 1 local column + .withReferredColumns( + Collections.singletonList(targetCol1FQN)); // 1 referred column (1:1 mapping) + + sourceTableV2.setTableConstraints(Collections.singletonList(foreignKeyConstraint)); + + Table updatedSourceTable = + tableTest.patchEntity(sourceTable.getId(), originalJson, sourceTableV2, ADMIN_AUTH_HEADERS); + + // Verify constraint was created correctly + assertNotNull(updatedSourceTable.getTableConstraints()); + assertEquals(1, updatedSourceTable.getTableConstraints().size()); + TableConstraint constraint = updatedSourceTable.getTableConstraints().getFirst(); + assertEquals(TableConstraint.ConstraintType.FOREIGN_KEY, constraint.getConstraintType()); + assertEquals(1, constraint.getColumns().size()); // 1 local column + assertEquals(1, constraint.getReferredColumns().size()); // 1 referred column (1:1 mapping) + + // Export recursively to CSV - this should include table constraints + String exportedCsv = exportCsvRecursive(schema.getFullyQualifiedName()); + assertNotNull(exportedCsv); + + List csvLines = List.of(exportedCsv.split(CsvUtil.LINE_SEPARATOR)); + assertTrue(csvLines.size() > 1, "Export should contain schema, tables, and columns"); + + // Modify CSV to update some metadata while preserving structure + String header = csvLines.getFirst(); + List modified = new ArrayList<>(); + modified.add(header); + + for (String line : csvLines.subList(1, csvLines.size())) { + if (line.contains("source_table") && line.contains("table")) { + // Update table description + line = line.replace("source_table", "source_table Updated via CSV import"); + } + modified.add(line); + } + + String newCsv = String.join(CsvUtil.LINE_SEPARATOR, modified) + CsvUtil.LINE_SEPARATOR; + + // Import the modified CSV recursively + CsvImportResult result = importCsvRecursive(schema.getFullyQualifiedName(), newCsv, false); + assertEquals(ApiStatus.SUCCESS, result.getStatus()); + + // Fetch the updated source table and verify constraints are preserved + Table importedSourceTable = + tableTest.getEntityByName( + updatedSourceTable.getFullyQualifiedName(), + "tableConstraints,columns,description", + ADMIN_AUTH_HEADERS); + + // Verify table constraints are still present after CSV import + assertNotNull( + importedSourceTable.getTableConstraints(), + "Table constraints should be preserved after CSV import"); + assertEquals( + 1, + importedSourceTable.getTableConstraints().size(), + "Should have exactly one table constraint"); + + TableConstraint preservedConstraint = importedSourceTable.getTableConstraints().getFirst(); + assertEquals( + TableConstraint.ConstraintType.FOREIGN_KEY, preservedConstraint.getConstraintType()); + assertEquals(1, preservedConstraint.getColumns().size(), "Should have 1 local column"); + assertEquals( + 1, + preservedConstraint.getReferredColumns().size(), + "Should have 1 referred column (1:1 mapping)"); + + // Verify the specific column references are preserved + assertEquals("user_ref", preservedConstraint.getColumns().getFirst()); + assertTrue( + preservedConstraint.getReferredColumns().contains(targetCol1FQN), + "Should contain target column FQN"); + + // Verify search index building works without crashing + assertDoesNotThrow( + () -> { + Entity.buildSearchIndex(Entity.TABLE, importedSourceTable); + }, + "Search index building should not crash with table constraints after CSV import"); + + // Verify target table is also intact + Table importedTargetTable = + tableTest.getEntityByName( + targetTable.getFullyQualifiedName(), "columns", ADMIN_AUTH_HEADERS); + + assertNotNull(importedTargetTable.getColumns()); + assertEquals( + 2, importedTargetTable.getColumns().size(), "Target table should still have 2 columns"); + } + @Override public DatabaseSchema validateGetWithDifferentFields(DatabaseSchema schema, boolean byName) throws HttpResponseException {