Fix Column Lineage Validation issue (#22308)

* Fix Column Lineage Validation issue

* Add test for validations

---------

Co-authored-by: Sriharsha Chintalapani <harshach@users.noreply.github.com>

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

@@ -871,6 +871,7 @@ public class LineageRepository {
         return result;
       }
       case API_ENDPOINT -> {
+        Set<String> result = new HashSet<>();
         APIEndpoint apiEndpoint =
             Entity.getEntity(
                 API_ENDPOINT,
@@ -878,15 +879,20 @@ public class LineageRepository {
                 "responseSchema,requestSchema",
                 Include.NON_DELETED);
         if (apiEndpoint.getResponseSchema() != null) {
-          return CommonUtil.getChildrenNames(
+          result.addAll(
-              apiEndpoint.getResponseSchema().getSchemaFields(),
+              CommonUtil.getChildrenNames(
+                  listOrEmpty(apiEndpoint.getResponseSchema().getSchemaFields()),
                   "getChildren",
-              apiEndpoint.getFullyQualifiedName());
+                  apiEndpoint.getFullyQualifiedName()));
         }
-        return CommonUtil.getChildrenNames(
+        if (apiEndpoint.getRequestSchema() != null) {
-            apiEndpoint.getRequestSchema().getSchemaFields(),
+          result.addAll(
+              CommonUtil.getChildrenNames(
+                  listOrEmpty(apiEndpoint.getRequestSchema().getSchemaFields()),
                   "getChildren",
-            apiEndpoint.getFullyQualifiedName());
+                  apiEndpoint.getFullyQualifiedName()));
+        }
+        return result;
       }
       case METRIC -> {
         LOG.info("Metric column level lineage is not supported");
@@ -896,8 +902,10 @@ public class LineageRepository {
         LOG.info("Pipeline column level lineage is not supported");
         return new HashSet<>();
       }
-      default -> throw new IllegalArgumentException(
+      default -> {
-          String.format("Unsupported Entity Type %s for lineage", entityReference.getType()));
+        LOG.error("Unsupported Entity Type {} for column lineage", entityReference.getType());
+        return new HashSet<>();
+      }
     }
   }
 

openmetadata-service/src/test/java/org/openmetadata/service/resources/lineage/LineageResourceTest.java

@@ -28,16 +28,19 @@ import static org.openmetadata.service.util.TestUtils.assertResponse;
 import jakarta.ws.rs.client.WebTarget;
 import jakarta.ws.rs.core.Response.Status;
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.net.URISyntaxException;
 import java.net.URLEncoder;
 import java.text.ParseException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.http.client.HttpResponseException;
@@ -71,18 +74,21 @@ import org.openmetadata.schema.tests.TestCase;
 import org.openmetadata.schema.tests.TestDefinition;
 import org.openmetadata.schema.tests.TestSuite;
 import org.openmetadata.schema.tests.type.TestCaseStatus;
+import org.openmetadata.schema.type.Column;
 import org.openmetadata.schema.type.ColumnLineage;
 import org.openmetadata.schema.type.ContainerDataModel;
 import org.openmetadata.schema.type.Edge;
 import org.openmetadata.schema.type.EntitiesEdge;
 import org.openmetadata.schema.type.EntityLineage;
 import org.openmetadata.schema.type.EntityReference;
+import org.openmetadata.schema.type.Field;
 import org.openmetadata.schema.type.LineageDetails;
 import org.openmetadata.schema.type.MetadataOperation;
 import org.openmetadata.schema.type.lineage.NodeInformation;
 import org.openmetadata.schema.utils.JsonUtils;
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.OpenMetadataApplicationTest;
+import org.openmetadata.service.jdbi3.LineageRepository;
 import org.openmetadata.service.resources.dashboards.DashboardResourceTest;
 import org.openmetadata.service.resources.databases.TableResourceTest;
 import org.openmetadata.service.resources.datamodels.DashboardDataModelResourceTest;
@@ -901,4 +907,158 @@ public class LineageResourceTest extends OpenMetadataApplicationTest {
       assertEdgeFromLineage(lineage.getDownstreamEdges(), expectedDownstreamEdge);
     }
   }
+
+  @Order(8)
+  @Test
+  void test_getChildrenNames_AllEntityTypes() throws Exception {
+    LineageRepository lineageRepository = new LineageRepository();
+    Method getChildrenNamesMethod =
+        LineageRepository.class.getDeclaredMethod("getChildrenNames", EntityReference.class);
+    getChildrenNamesMethod.setAccessible(true);
+
+    // Test Table Entity - should return column children
+    EntityReference tableRef = TABLES.get(0).getEntityReference();
+    Set<String> tableChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, tableRef);
+    assertFalse(tableChildren.isEmpty(), "Table should have column children");
+    assertTrue(tableChildren.size() >= 3, "Table should have at least 3 columns");
+    Set<String> expectedColumns =
+        TABLES.get(0).getColumns().stream().map(Column::getName).collect(Collectors.toSet());
+    assertTrue(
+        tableChildren.containsAll(expectedColumns),
+        "Table children should contain expected column names: " + expectedColumns);
+
+    // Test Topic Entity - should return schema field children
+    EntityReference topicRef = TOPIC.getEntityReference();
+    Set<String> topicChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, topicRef);
+    assertFalse(topicChildren.isEmpty(), "Topic should have schema field children");
+    assertTrue(topicChildren.size() >= 1, "Topic should have at least 1 schema field");
+    Set<String> expectedFields =
+        TOPIC.getMessageSchema().getSchemaFields().stream()
+            .map(Field::getName)
+            .sorted()
+            .collect(Collectors.toCollection(LinkedHashSet::new));
+    assertEquals(
+        expectedFields,
+        topicChildren.stream().sorted().collect(Collectors.toCollection(LinkedHashSet::new)),
+        "Topic children should contain expected field names: " + expectedFields);
+
+    // Test Container Entity - should return data model column children
+    EntityReference containerRef = CONTAINER.getEntityReference();
+    Set<String> containerChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, containerRef);
+    assertFalse(containerChildren.isEmpty(), "Container should have data model column children");
+    assertTrue(containerChildren.size() >= 2, "Container should have at least 2 columns");
+    Set<String> expectedContainerField =
+        CONTAINER.getDataModel().getColumns().stream()
+            .map(Column::getName)
+            .sorted()
+            .collect(Collectors.toCollection(LinkedHashSet::new));
+    assertEquals(
+        expectedContainerField,
+        containerChildren.stream().sorted().collect(Collectors.toCollection(LinkedHashSet::new)),
+        "Container children should contain expected column names: " + expectedContainerField);
+
+    // Test DashboardDataModel Entity - should return column children
+    EntityReference dataModelRef = DATA_MODEL.getEntityReference();
+    Set<String> dataModelChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, dataModelRef);
+    assertFalse(dataModelChildren.isEmpty(), "DashboardDataModel should have column children");
+    assertTrue(dataModelChildren.size() >= 3, "DashboardDataModel should have at least 3 columns");
+    Set<String> expectedDataModelColumns =
+        DATA_MODEL.getColumns().stream()
+            .map(Column::getName)
+            .sorted()
+            .collect(Collectors.toCollection(LinkedHashSet::new));
+    assertTrue(
+        dataModelChildren.stream().sorted().toList().containsAll(expectedDataModelColumns),
+        "DashboardDataModel children should contain expected column names: "
+            + expectedDataModelColumns);
+
+    // Test Dashboard Entity - should return chart children without FQN prefix
+    EntityReference dashboardRef = DASHBOARD.getEntityReference();
+    Set<String> dashboardChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, dashboardRef);
+    assertFalse(dashboardChildren.isEmpty(), "Dashboard should have chart children");
+    assertTrue(dashboardChildren.size() >= 2, "Dashboard should have at least 2 charts");
+    Set<String> expectedChartNames =
+        DASHBOARD.getCharts().stream()
+            .map(
+                chart ->
+                    chart
+                        .getFullyQualifiedName()
+                        .replace(DASHBOARD.getFullyQualifiedName() + ".", ""))
+            .collect(Collectors.toSet());
+    assertEquals(
+        expectedChartNames,
+        dashboardChildren,
+        "Dashboard children should match expected chart names without FQN prefix");
+    for (String chartName : dashboardChildren) {
+      assertFalse(
+          chartName.contains(DASHBOARD.getFullyQualifiedName() + "."),
+          "Chart name should not contain dashboard FQN prefix: " + chartName);
+    }
+
+    // Test MlModel Entity - should return feature children without FQN prefix
+    EntityReference mlModelRef = ML_MODEL.getEntityReference();
+    Set<String> mlModelChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, mlModelRef);
+    assertFalse(mlModelChildren.isEmpty(), "MlModel should have ML feature children");
+    assertTrue(mlModelChildren.size() >= 2, "MlModel should have at least 2 ML features");
+    Set<String> expectedFeatureNames =
+        ML_MODEL.getMlFeatures().stream()
+            .map(
+                feature ->
+                    feature
+                        .getFullyQualifiedName()
+                        .replace(ML_MODEL.getFullyQualifiedName() + ".", ""))
+            .collect(Collectors.toSet());
+    assertEquals(
+        expectedFeatureNames,
+        mlModelChildren,
+        "MlModel children should match expected feature names without FQN prefix");
+    for (String featureName : mlModelChildren) {
+      assertFalse(
+          featureName.contains(ML_MODEL.getFullyQualifiedName() + "."),
+          "Feature name should not contain ML model FQN prefix: " + featureName);
+    }
+
+    // Test Topic Entity without schema - should return empty set
+    TopicResourceTest topicResourceTest = new TopicResourceTest();
+    CreateTopic topicRequest = topicResourceTest.createRequest("topicWithoutSchema");
+    topicRequest.setMessageSchema(null);
+    Topic topicWithoutSchema = topicResourceTest.createEntity(topicRequest, ADMIN_AUTH_HEADERS);
+    EntityReference topicWithoutSchemaRef = topicWithoutSchema.getEntityReference();
+    Set<String> topicWithoutSchemaChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, topicWithoutSchemaRef);
+    assertTrue(
+        topicWithoutSchemaChildren.isEmpty(),
+        "Topic without message schema should return empty set");
+
+    // Test Container Entity without data model - should return empty set
+    ContainerResourceTest containerResourceTest = new ContainerResourceTest();
+    CreateContainer containerRequest =
+        containerResourceTest.createRequest("containerWithoutDataModel");
+    containerRequest.setDataModel(null);
+    Container containerWithoutDataModel =
+        containerResourceTest.createEntity(containerRequest, ADMIN_AUTH_HEADERS);
+    EntityReference containerWithoutDataModelRef = containerWithoutDataModel.getEntityReference();
+    Set<String> containerWithoutDataModelChildren =
+        (Set<String>)
+            getChildrenNamesMethod.invoke(lineageRepository, containerWithoutDataModelRef);
+    assertTrue(
+        containerWithoutDataModelChildren.isEmpty(),
+        "Container without data model should return empty set");
+
+    // Test unknown entity type - should return empty set
+    EntityReference unknownRef =
+        new EntityReference()
+            .withId(UUID.randomUUID())
+            .withType("UNKNOWN_TYPE")
+            .withFullyQualifiedName("test.unknown");
+    Set<String> unknownChildren =
+        (Set<String>) getChildrenNamesMethod.invoke(lineageRepository, unknownRef);
+    assertTrue(unknownChildren.isEmpty(), "Unknown entity type should return empty set");
+  }
 }