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!

(cherry picked from commit 8d7db583049d8f93e0c1330255a2c43588f6eaff)

openmetadata-service/src/main/java/org/openmetadata/csv/EntityCsv.java

@@ -1069,8 +1069,11 @@ public abstract class EntityCsv<T extends EntityInterface> {
     Database database;
     try {
       database =
-          Entity.getEntityByName(
+          Entity.getEntityByNameWithExcludedFields(
-              DATABASE, dbFQN, "name,displayName,fullyQualifiedName,service", Include.NON_DELETED);
+              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<T extends EntityInterface> {
     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<T extends EntityInterface> {
 
     try {
       table =
-          Entity.getEntityByName(
+          Entity.getEntityByNameWithExcludedFields(
-              TABLE, tableFqn, "name,displayName,fullyQualifiedName,columns", Include.NON_DELETED);
+              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
 

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> 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> 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 <T> List<T> getEntityByNames(
       String entityType, List<String> tagFQNs, String fields, Include include) {
     EntityRepository<?> entityRepository = Entity.getEntityRepository(entityType);

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java

@@ -787,6 +787,19 @@ public abstract class EntityRepository<T extends EntityInterface> {
     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);

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<String> allowedFields, Set<String> excludeFields) {
+      Set<String> resultFields = new HashSet<>(allowedFields);
+      if (excludeFields != null) {
+        resultFields.removeAll(excludeFields);
+      }
+      return new Fields(allowedFields, resultFields);
+    }
+
+    public static Fields createWithExcludedFields(
+        Set<String> allowedFields, String excludeFieldsParam) {
+      Set<String> 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);

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<String> 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<String> 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 {