Fix#10584: Update Lineage API to support lineage from table to dashboard data model (#11489)

* Fix: Update Lineage API to support lineage from table to dashboard data model

* Minor change

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

@@ -265,7 +265,7 @@ public interface CollectionDAO {
   WorkflowDAO workflowDAO();
 
   @CreateSqlObject
-  DataModelDAO dataModelDAO();
+  DataModelDAO dashboardDataModelDAO();
 
   interface DashboardDAO extends EntityDAO<Dashboard> {
     @Override

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

@@ -52,7 +52,7 @@ public class DashboardDataModelRepository extends EntityRepository<DashboardData
         DashboardDataModelResource.COLLECTION_PATH,
         Entity.DASHBOARD_DATA_MODEL,
         DashboardDataModel.class,
-        dao.dataModelDAO(),
+        dao.dashboardDataModelDAO(),
         dao,
         DATA_MODEL_PATCH_FIELDS,
         DATA_MODEL_UPDATE_FIELDS,

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

@@ -19,6 +19,7 @@ import java.util.List;
 import java.util.UUID;
 import java.util.stream.Collectors;
 import org.jdbi.v3.sqlobject.transaction.Transaction;
+import org.openmetadata.schema.ColumnsEntityInterface;
 import org.openmetadata.schema.api.lineage.AddLineage;
 import org.openmetadata.schema.entity.data.Table;
 import org.openmetadata.schema.type.ColumnLineage;
@@ -89,12 +90,13 @@ public class LineageRepository {
     }
 
     List<ColumnLineage> columnsLineage = details.getColumnsLineage();
-    if (!from.getType().equals(Entity.TABLE) || !to.getType().equals(Entity.TABLE)) {
+    if (areValidEntities(from, to)) {
-      throw new IllegalArgumentException("Column level lineage is only allowed between two tables.");
+      throw new IllegalArgumentException(
+          "Column level lineage is only allowed between two tables or from table to dashboard.");
     }
 
     Table fromTable = dao.tableDAO().findEntityById(from.getId());
-    Table toTable = dao.tableDAO().findEntityById(to.getId());
+    ColumnsEntityInterface toTable = getToEntity(to);
     if (columnsLineage != null) {
       for (ColumnLineage columnLineage : columnsLineage) {
         for (String fromColumn : columnLineage.getFromColumns()) {
@@ -112,6 +114,17 @@ public class LineageRepository {
     return JsonUtils.pojoToJson(details);
   }
 
+  private ColumnsEntityInterface getToEntity(EntityReference from) throws IOException {
+    return from.getType().equals(Entity.TABLE)
+        ? dao.tableDAO().findEntityById(from.getId())
+        : dao.dashboardDataModelDAO().findEntityById(from.getId());
+  }
+
+  private boolean areValidEntities(EntityReference from, EntityReference to) {
+    return !from.getType().equals(Entity.TABLE)
+        || !(to.getType().equals(Entity.TABLE) || to.getType().equals(Entity.DASHBOARD_DATA_MODEL));
+  }
+
   @Transaction
   public boolean deleteLineage(String fromEntity, String fromId, String toEntity, String toId) throws IOException {
     // Validate from entity

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

@@ -41,8 +41,11 @@ import org.junit.jupiter.api.Order;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInfo;
 import org.junit.jupiter.api.TestMethodOrder;
+import org.openmetadata.schema.EntityInterface;
+import org.openmetadata.schema.api.data.CreateDashboardDataModel;
 import org.openmetadata.schema.api.data.CreateTable;
 import org.openmetadata.schema.api.lineage.AddLineage;
+import org.openmetadata.schema.entity.data.DashboardDataModel;
 import org.openmetadata.schema.entity.data.Table;
 import org.openmetadata.schema.entity.teams.Role;
 import org.openmetadata.schema.entity.teams.User;
@@ -56,6 +59,7 @@ import org.openmetadata.schema.type.MetadataOperation;
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.OpenMetadataApplicationTest;
 import org.openmetadata.service.resources.databases.TableResourceTest;
+import org.openmetadata.service.resources.datamodels.DashboardDataModelResourceTest;
 import org.openmetadata.service.resources.teams.RoleResource;
 import org.openmetadata.service.resources.teams.RoleResourceTest;
 import org.openmetadata.service.resources.teams.UserResourceTest;
@@ -66,9 +70,12 @@ import org.openmetadata.service.util.TestUtils;
 public class LineageResourceTest extends OpenMetadataApplicationTest {
   public static final List<Table> TABLES = new ArrayList<>();
   public static final int TABLE_COUNT = 10;
-
   private static final String DATA_STEWARD_ROLE_NAME = "DataSteward";
 
+  private static DashboardDataModel DATA_MODEL;
+
+  private static Table TABLE_DATA_MODEL_LINEAGE;
+
   @BeforeAll
   public static void setup(TestInfo test) throws IOException, URISyntaxException {
     // Create TABLE_COUNT number of tables
@@ -78,6 +85,14 @@ public class LineageResourceTest extends OpenMetadataApplicationTest {
       CreateTable createTable = tableResourceTest.createRequest(test, i);
       TABLES.add(tableResourceTest.createEntity(createTable, ADMIN_AUTH_HEADERS));
     }
+
+    // Entities to test lineage DashboardDataModel <-> Table
+    DashboardDataModelResourceTest dashboardResourceTest = new DashboardDataModelResourceTest();
+    CreateDashboardDataModel createDashboardDataModel = dashboardResourceTest.createRequest(test);
+    DATA_MODEL = dashboardResourceTest.createEntity(createDashboardDataModel, ADMIN_AUTH_HEADERS);
+    CreateTable createTable = tableResourceTest.createRequest(test, TABLE_COUNT);
+    createTable.setColumns(createDashboardDataModel.getColumns());
+    TABLE_DATA_MODEL_LINEAGE = tableResourceTest.createEntity(createTable, ADMIN_AUTH_HEADERS);
   }
 
   @Order(1)
@@ -286,6 +301,32 @@ public class LineageResourceTest extends OpenMetadataApplicationTest {
     addEdge(TABLES.get(0), TABLES.get(1), details, ADMIN_AUTH_HEADERS);
   }
 
+  @Order(4)
+  @Test
+  void putLineageFromDashboardDataModelToTable() throws HttpResponseException {
+
+    // Add column lineage dashboard.d1 -> table.c1
+    LineageDetails details = new LineageDetails();
+    String d1c1FQN = DATA_MODEL.getColumns().get(0).getFullyQualifiedName();
+    String d1c2FQN = DATA_MODEL.getColumns().get(1).getFullyQualifiedName();
+    String d1c3FQN = DATA_MODEL.getColumns().get(2).getFullyQualifiedName();
+    String c1c1FQN = TABLE_DATA_MODEL_LINEAGE.getColumns().get(0).getFullyQualifiedName();
+    String c1c2FQN = TABLE_DATA_MODEL_LINEAGE.getColumns().get(1).getFullyQualifiedName();
+    String c1c3FQN = TABLE_DATA_MODEL_LINEAGE.getColumns().get(2).getFullyQualifiedName();
+
+    List<ColumnLineage> lineage = details.getColumnsLineage();
+    lineage.add(new ColumnLineage().withFromColumns(List.of(c1c1FQN)).withToColumn(d1c1FQN));
+    lineage.add(new ColumnLineage().withFromColumns(List.of(c1c2FQN)).withToColumn(d1c2FQN));
+    lineage.add(new ColumnLineage().withFromColumns(List.of(c1c3FQN)).withToColumn(d1c3FQN));
+
+    addEdge(TABLE_DATA_MODEL_LINEAGE, DATA_MODEL, details, ADMIN_AUTH_HEADERS);
+
+    assertResponse(
+        () -> addEdge(DATA_MODEL, TABLE_DATA_MODEL_LINEAGE, details, ADMIN_AUTH_HEADERS),
+        BAD_REQUEST,
+        "Column level lineage is only allowed between two tables or from table to dashboard.");
+  }
+
   public Edge getEdge(Table from, Table to) {
     return getEdge(from.getId(), to.getId(), null);
   }
@@ -298,7 +339,8 @@ public class LineageResourceTest extends OpenMetadataApplicationTest {
     addEdge(from, to, null, ADMIN_AUTH_HEADERS);
   }
 
-  private void addEdge(Table from, Table to, LineageDetails details, Map<String, String> authHeaders)
+  private void addEdge(
+      EntityInterface from, EntityInterface to, LineageDetails details, Map<String, String> authHeaders)
       throws HttpResponseException {
     if (details != null) {
       details.setSqlQuery("select *;");

openmetadata-spec/src/main/java/org/openmetadata/schema/ColumnsEntityInterface.java

@@ -0,0 +1,28 @@
+/*
+ *  Copyright 2022 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.schema;
+
+import java.util.List;
+import org.openmetadata.schema.type.Column;
+
+/**
+ * Interface to be implemented by entities with a list of Column and FullyQualifiedName. It is used when adding lineage
+ * between different entities.
+ */
+public interface ColumnsEntityInterface {
+
+  String getFullyQualifiedName();
+
+  List<Column> getColumns();
+}

openmetadata-spec/src/main/resources/json/schema/entity/data/dashboardDataModel.json

@@ -6,7 +6,7 @@
   "description": "Dashboard Data Model entity definition. Data models are the schemas used to build dashboards, charts, or other data assets.",
   "type": "object",
   "javaType": "org.openmetadata.schema.entity.data.DashboardDataModel",
-  "javaInterfaces": ["org.openmetadata.schema.EntityInterface"],
+  "javaInterfaces": ["org.openmetadata.schema.EntityInterface", "org.openmetadata.schema.ColumnsEntityInterface"],
   "definitions": {
     "dataModelType": {
       "javaType": "org.openmetadata.schema.type.DataModelType",

openmetadata-spec/src/main/resources/json/schema/entity/data/table.json

@@ -7,7 +7,7 @@
   "type": "object",
   "javaType": "org.openmetadata.schema.entity.data.Table",
   "javaInterfaces": [
-    "org.openmetadata.schema.EntityInterface"
+    "org.openmetadata.schema.EntityInterface", "org.openmetadata.schema.ColumnsEntityInterface"
   ],
   "definitions": {
     "entityName": {