Fix #2555: PUT APIs should only allow admin and owners to update a resource in secure env (#2608)

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/AirflowPipelineRepository.java

@@ -53,7 +53,9 @@ public class AirflowPipelineRepository extends EntityRepository<AirflowPipeline>
   }
 
   public static String getFQN(AirflowPipeline airflowPipeline) {
-    return (airflowPipeline.getService().getName() + "." + airflowPipeline.getName());
+    return (airflowPipeline != null && airflowPipeline.getService() != null)
+        ? (airflowPipeline.getService().getName() + "." + airflowPipeline.getName())
+        : null;
   }
 
   @Transaction
@@ -162,7 +164,9 @@ public class AirflowPipelineRepository extends EntityRepository<AirflowPipeline>
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null
+          ? entity.getFullyQualifiedName()
+          : AirflowPipelineRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/ChartRepository.java

@@ -52,7 +52,9 @@ public class ChartRepository extends EntityRepository<Chart> {
   }
 
   public static String getFQN(Chart chart) {
-    return (chart.getService().getName() + "." + chart.getName());
+    return (chart != null && chart.getService() != null)
+        ? (chart.getService().getName() + "." + chart.getName())
+        : null;
   }
 
   @Override
@@ -158,7 +160,7 @@ public class ChartRepository extends EntityRepository<Chart> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null ? entity.getFullyQualifiedName() : ChartRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/DashboardRepository.java

@@ -58,7 +58,9 @@ public class DashboardRepository extends EntityRepository<Dashboard> {
   }
 
   public static String getFQN(Dashboard dashboard) {
-    return (dashboard.getService().getName() + "." + dashboard.getName());
+    return (dashboard != null && dashboard.getService() != null)
+        ? (dashboard.getService().getName() + "." + dashboard.getName())
+        : null;
   }
 
   @Override
@@ -271,7 +273,9 @@ public class DashboardRepository extends EntityRepository<Dashboard> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null
+          ? entity.getFullyQualifiedName()
+          : DashboardRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/DatabaseRepository.java

@@ -58,7 +58,9 @@ public class DatabaseRepository extends EntityRepository<Database> {
   }
 
   public static String getFQN(Database database) {
-    return (database.getService().getName() + "." + database.getName());
+    return (database != null && database.getService() != null)
+        ? (database.getService().getName() + "." + database.getName())
+        : null;
   }
 
   @Transaction
@@ -243,7 +245,9 @@ public class DatabaseRepository extends EntityRepository<Database> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null
+          ? entity.getFullyQualifiedName()
+          : DatabaseRepository.getFQN(entity);
     }
 
     @Override

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

@@ -30,6 +30,7 @@ import java.net.URI;
 import java.security.GeneralSecurityException;
 import java.text.ParseException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -559,6 +560,23 @@ public abstract class EntityRepository<T> {
     }
   }
 
+  public final EntityReference getOriginalOwner(T entity) throws IOException, ParseException {
+    final String FIELDS = "owner";
+    final List<String> FIELD_LIST = Arrays.asList(FIELDS.replace(" ", "").split(","));
+    EntityUtil.Fields fields = new EntityUtil.Fields(FIELD_LIST, FIELDS);
+    EntityReference owner = null;
+    // Try to find the owner if entity exists
+    try {
+      String fqn = getFullyQualifiedName(entity);
+      entity = getByName(null, fqn, fields);
+      owner = helper(entity).validateOwnerOrNull();
+    } catch (EntityNotFoundException e) {
+      // If entity is not found, we can return null for owner and ignore
+      // this exception
+    }
+    return owner;
+  }
+
   protected EntityReference getOwner(T entity) throws IOException, ParseException {
     return helper(entity).getOwnerOrNull();
   }

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/LocationRepository.java

@@ -156,7 +156,9 @@ public class LocationRepository extends EntityRepository<Location> {
   }
 
   public static String getFQN(Location location) {
-    return (location.getService().getName() + "." + location.getName());
+    return (location != null && location.getService() != null)
+        ? (location.getService().getName() + "." + location.getName())
+        : null;
   }
 
   @Transaction
@@ -281,7 +283,9 @@ public class LocationRepository extends EntityRepository<Location> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null
+          ? entity.getFullyQualifiedName()
+          : LocationRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/MlModelRepository.java

@@ -278,7 +278,7 @@ public class MlModelRepository extends EntityRepository<MlModel> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null ? entity.getFullyQualifiedName() : MlModelRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/PipelineRepository.java

@@ -62,7 +62,9 @@ public class PipelineRepository extends EntityRepository<Pipeline> {
   }
 
   public static String getFQN(Pipeline pipeline) {
-    return (pipeline.getService().getName() + "." + pipeline.getName());
+    return (pipeline != null && pipeline.getService() != null)
+        ? (pipeline.getService().getName() + "." + pipeline.getName())
+        : null;
   }
 
   @Transaction
@@ -200,7 +202,9 @@ public class PipelineRepository extends EntityRepository<Pipeline> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null
+          ? entity.getFullyQualifiedName()
+          : PipelineRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/PolicyRepository.java

@@ -309,7 +309,7 @@ public class PolicyRepository extends EntityRepository<Policy> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null ? entity.getFullyQualifiedName() : PolicyRepository.getFQN(entity);
     }
 
     public List<Object> getRules() {

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/TableRepository.java

@@ -127,7 +127,9 @@ public class TableRepository extends EntityRepository<Table> {
   }
 
   public static String getFQN(Table table) {
-    return (table.getDatabase().getName() + "." + table.getName());
+    return (table != null && table.getDatabase() != null)
+        ? (table.getDatabase().getName() + "." + table.getName())
+        : null;
   }
 
   @Transaction
@@ -667,7 +669,7 @@ public class TableRepository extends EntityRepository<Table> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null ? entity.getFullyQualifiedName() : TableRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/TopicRepository.java

@@ -41,7 +41,9 @@ public class TopicRepository extends EntityRepository<Topic> {
   private static final Fields TOPIC_PATCH_FIELDS = new Fields(TopicResource.FIELD_LIST, "owner,tags");
 
   public static String getFQN(Topic topic) {
-    return (topic.getService().getName() + "." + topic.getName());
+    return (topic != null && topic.getService() != null)
+        ? (topic.getService().getName() + "." + topic.getName())
+        : null;
   }
 
   public TopicRepository(CollectionDAO dao) {
@@ -172,7 +174,7 @@ public class TopicRepository extends EntityRepository<Topic> {
 
     @Override
     public String getFullyQualifiedName() {
-      return entity.getFullyQualifiedName();
+      return entity.getFullyQualifiedName() != null ? entity.getFullyQualifiedName() : TopicRepository.getFQN(entity);
     }
 
     @Override

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/charts/ChartResource.java

@@ -350,8 +350,8 @@ public class ChartResource {
   public Response createOrUpdate(
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateChart create)
       throws IOException, ParseException {
-
     Chart chart = getChart(securityContext, create);
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(chart));
     PutResponse<Chart> response = dao.createOrUpdate(uriInfo, chart);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/dashboards/DashboardResource.java

@@ -360,6 +360,7 @@ public class DashboardResource {
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateDashboard create)
       throws IOException, ParseException {
     Dashboard dashboard = getDashboard(securityContext, create);
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(dashboard));
     PutResponse<Dashboard> response = dao.createOrUpdate(uriInfo, dashboard);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/databases/DatabaseResource.java

@@ -370,6 +370,7 @@ public class DatabaseResource {
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateDatabase create)
       throws IOException, ParseException {
     Database database = getDatabase(securityContext, create);
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(database));
     PutResponse<Database> response = dao.createOrUpdate(uriInfo, database);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/databases/TableResource.java

@@ -13,8 +13,6 @@
 
 package org.openmetadata.catalog.resources.databases;
 
-import static org.openmetadata.catalog.Entity.helper;
-
 import io.swagger.annotations.Api;
 import io.swagger.v3.oas.annotations.ExternalDocumentation;
 import io.swagger.v3.oas.annotations.Operation;
@@ -61,7 +59,6 @@ import org.openmetadata.catalog.security.Authorizer;
 import org.openmetadata.catalog.security.SecurityUtil;
 import org.openmetadata.catalog.type.DataModel;
 import org.openmetadata.catalog.type.EntityHistory;
-import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.type.SQLQuery;
 import org.openmetadata.catalog.type.TableData;
@@ -322,7 +319,7 @@ public class TableResource {
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateTable create)
       throws IOException, ParseException {
     Table table = getTable(securityContext, create);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, helper(table).validateOwnerOrNull());
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(table));
     PutResponse<Table> response = dao.createOrUpdate(uriInfo, validateNewTable(table));
     addHref(uriInfo, response.getEntity());
     return response.toResponse();
@@ -559,6 +556,6 @@ public class TableResource {
         .withUpdatedBy(securityContext.getUserPrincipal().getName())
         .withOwner(create.getOwner())
         .withUpdatedAt(System.currentTimeMillis())
-        .withDatabase(new EntityReference().withId(create.getDatabase()).withType(Entity.DATABASE));
+        .withDatabase(create.getDatabase());
   }
 }

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/locations/LocationResource.java

@@ -371,7 +371,7 @@ public class LocationResource {
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateLocation create)
       throws IOException, ParseException {
     Location location = getLocation(securityContext, create);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOwnerReference(location));
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(location));
     PutResponse<Location> response = dao.createOrUpdate(uriInfo, validateNewLocation(location));
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/mlmodels/MlModelResource.java

@@ -298,7 +298,7 @@ public class MlModelResource {
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateMlModel create)
       throws IOException, ParseException {
     MlModel mlModel = getMlModel(securityContext, create);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOwnerReference(mlModel));
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(mlModel));
     PutResponse<MlModel> response = dao.createOrUpdate(uriInfo, mlModel);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/operations/AirflowPipelineResource.java

@@ -388,11 +388,12 @@ public class AirflowPipelineResource {
         @ApiResponse(responseCode = "400", description = "Bad request")
       })
   public Response createOrUpdate(
-      @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateAirflowPipeline create)
+      @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateAirflowPipeline update)
       throws IOException, ParseException {
-    AirflowPipeline airflowPipeline = getAirflowPipeline(securityContext, create);
+    AirflowPipeline pipeline = getAirflowPipeline(securityContext, update);
-    PutResponse<AirflowPipeline> response = dao.createOrUpdate(uriInfo, airflowPipeline);
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(pipeline));
-    deploy(airflowPipeline);
+    PutResponse<AirflowPipeline> response = dao.createOrUpdate(uriInfo, pipeline);
+    deploy(pipeline);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();
   }

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/pipelines/PipelineResource.java

@@ -359,6 +359,7 @@ public class PipelineResource {
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreatePipeline create)
       throws IOException, ParseException {
     Pipeline pipeline = getPipeline(securityContext, create);
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(pipeline));
     PutResponse<Pipeline> response = dao.createOrUpdate(uriInfo, pipeline);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/policies/PolicyResource.java

@@ -360,6 +360,7 @@ public class PolicyResource {
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreatePolicy create)
       throws IOException, ParseException {
     Policy policy = getPolicy(securityContext, create);
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(policy));
     PutResponse<Policy> response = dao.createOrUpdate(uriInfo, policy);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/services/dashboard/DashboardServiceResource.java

@@ -13,8 +13,6 @@
 
 package org.openmetadata.catalog.resources.services.dashboard;
 
-import static org.openmetadata.catalog.Entity.helper;
-
 import io.swagger.annotations.Api;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
@@ -52,14 +50,12 @@ import javax.ws.rs.core.UriInfo;
 import org.openmetadata.catalog.Entity;
 import org.openmetadata.catalog.api.services.CreateDashboardService;
 import org.openmetadata.catalog.entity.services.DashboardService;
-import org.openmetadata.catalog.exception.EntityNotFoundException;
 import org.openmetadata.catalog.jdbi3.CollectionDAO;
 import org.openmetadata.catalog.jdbi3.DashboardServiceRepository;
 import org.openmetadata.catalog.resources.Collection;
 import org.openmetadata.catalog.security.Authorizer;
 import org.openmetadata.catalog.security.SecurityUtil;
 import org.openmetadata.catalog.type.EntityHistory;
-import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.util.EntityUtil;
 import org.openmetadata.catalog.util.RestUtil;
@@ -316,18 +312,8 @@ public class DashboardServiceResource {
   public Response update(
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateDashboardService update)
       throws IOException, ParseException {
-    EntityUtil.Fields fields = new EntityUtil.Fields(FIELD_LIST, FIELDS);
+    DashboardService service = getService(update, securityContext);
-    EntityReference owner = null;
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(service));
-    DashboardService service;
-    // Try to find the owner if entity exists
-    try {
-      service = dao.getByName(uriInfo, update.getName(), fields);
-      owner = helper(service).validateOwnerOrNull();
-    } catch (EntityNotFoundException e) {
-      // This is a create request if entity is not found. ignore exception
-    }
-    service = getService(update, securityContext);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, owner);
     PutResponse<DashboardService> response = dao.createOrUpdate(uriInfo, service, true);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/services/database/DatabaseServiceResource.java

@@ -13,8 +13,6 @@
 
 package org.openmetadata.catalog.resources.services.database;
 
-import static org.openmetadata.catalog.Entity.helper;
-
 import io.swagger.annotations.Api;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
@@ -52,14 +50,12 @@ import javax.ws.rs.core.UriInfo;
 import org.openmetadata.catalog.Entity;
 import org.openmetadata.catalog.api.services.CreateDatabaseService;
 import org.openmetadata.catalog.entity.services.DatabaseService;
-import org.openmetadata.catalog.exception.EntityNotFoundException;
 import org.openmetadata.catalog.jdbi3.CollectionDAO;
 import org.openmetadata.catalog.jdbi3.DatabaseServiceRepository;
 import org.openmetadata.catalog.resources.Collection;
 import org.openmetadata.catalog.security.Authorizer;
 import org.openmetadata.catalog.security.SecurityUtil;
 import org.openmetadata.catalog.type.EntityHistory;
-import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.util.EntityUtil;
 import org.openmetadata.catalog.util.RestUtil;
@@ -312,18 +308,8 @@ public class DatabaseServiceResource {
   public Response update(
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateDatabaseService update)
       throws IOException, ParseException {
-    EntityUtil.Fields fields = new EntityUtil.Fields(FIELD_LIST, FIELDS);
+    DatabaseService service = getService(update, securityContext);
-    EntityReference owner = null;
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(service));
-    DatabaseService service;
-    // Try to find the owner if entity exists
-    try {
-      service = dao.getByName(uriInfo, update.getName(), fields);
-      owner = helper(service).validateOwnerOrNull();
-    } catch (EntityNotFoundException e) {
-      // This is a create request if entity is not found. ignore exception
-    }
-    service = getService(update, securityContext);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, owner);
     PutResponse<DatabaseService> response = dao.createOrUpdate(uriInfo, service, true);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/services/messaging/MessagingServiceResource.java

@@ -13,8 +13,6 @@
 
 package org.openmetadata.catalog.resources.services.messaging;
 
-import static org.openmetadata.catalog.Entity.helper;
-
 import io.swagger.annotations.Api;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
@@ -52,14 +50,12 @@ import javax.ws.rs.core.UriInfo;
 import org.openmetadata.catalog.Entity;
 import org.openmetadata.catalog.api.services.CreateMessagingService;
 import org.openmetadata.catalog.entity.services.MessagingService;
-import org.openmetadata.catalog.exception.EntityNotFoundException;
 import org.openmetadata.catalog.jdbi3.CollectionDAO;
 import org.openmetadata.catalog.jdbi3.MessagingServiceRepository;
 import org.openmetadata.catalog.resources.Collection;
 import org.openmetadata.catalog.security.Authorizer;
 import org.openmetadata.catalog.security.SecurityUtil;
 import org.openmetadata.catalog.type.EntityHistory;
-import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.util.EntityUtil;
 import org.openmetadata.catalog.util.RestUtil;
@@ -323,18 +319,8 @@ public class MessagingServiceResource {
           String id,
       @Valid CreateMessagingService update)
       throws IOException, ParseException {
-    EntityUtil.Fields fields = new EntityUtil.Fields(FIELD_LIST, FIELDS);
+    MessagingService service = getService(update, securityContext);
-    EntityReference owner = null;
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(service));
-    MessagingService service;
-    // Try to find the owner if entity exists
-    try {
-      service = dao.getByName(uriInfo, update.getName(), fields);
-      owner = helper(service).validateOwnerOrNull();
-    } catch (EntityNotFoundException e) {
-      // This is a create request if entity is not found. ignore exception
-    }
-    service = getService(update, securityContext);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, owner);
     PutResponse<MessagingService> response = dao.createOrUpdate(uriInfo, service, true);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/services/pipeline/PipelineServiceResource.java

@@ -13,8 +13,6 @@
 
 package org.openmetadata.catalog.resources.services.pipeline;
 
-import static org.openmetadata.catalog.Entity.helper;
-
 import io.swagger.annotations.Api;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
@@ -52,14 +50,12 @@ import javax.ws.rs.core.UriInfo;
 import org.openmetadata.catalog.Entity;
 import org.openmetadata.catalog.api.services.CreatePipelineService;
 import org.openmetadata.catalog.entity.services.PipelineService;
-import org.openmetadata.catalog.exception.EntityNotFoundException;
 import org.openmetadata.catalog.jdbi3.CollectionDAO;
 import org.openmetadata.catalog.jdbi3.PipelineServiceRepository;
 import org.openmetadata.catalog.resources.Collection;
 import org.openmetadata.catalog.security.Authorizer;
 import org.openmetadata.catalog.security.SecurityUtil;
 import org.openmetadata.catalog.type.EntityHistory;
-import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.util.EntityUtil;
 import org.openmetadata.catalog.util.RestUtil;
@@ -319,18 +315,8 @@ public class PipelineServiceResource {
   public Response update(
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreatePipelineService update)
       throws IOException, ParseException {
-    EntityUtil.Fields fields = new EntityUtil.Fields(FIELD_LIST, FIELDS);
+    PipelineService service = getService(update, securityContext);
-    EntityReference owner = null;
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(service));
-    PipelineService service;
-    // Try to find the owner if entity exists
-    try {
-      service = dao.getByName(uriInfo, update.getName(), fields);
-      owner = helper(service).validateOwnerOrNull();
-    } catch (EntityNotFoundException e) {
-      // This is a create request if entity is not found. ignore exception
-    }
-    service = getService(update, securityContext);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, owner);
     PutResponse<PipelineService> response = dao.createOrUpdate(uriInfo, service, true);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/services/storage/StorageServiceResource.java

@@ -13,8 +13,6 @@
 
 package org.openmetadata.catalog.resources.services.storage;
 
-import static org.openmetadata.catalog.Entity.helper;
-
 import io.swagger.annotations.Api;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
@@ -52,14 +50,12 @@ import javax.ws.rs.core.UriInfo;
 import org.openmetadata.catalog.Entity;
 import org.openmetadata.catalog.api.services.CreateStorageService;
 import org.openmetadata.catalog.entity.services.StorageService;
-import org.openmetadata.catalog.exception.EntityNotFoundException;
 import org.openmetadata.catalog.jdbi3.CollectionDAO;
 import org.openmetadata.catalog.jdbi3.StorageServiceRepository;
 import org.openmetadata.catalog.resources.Collection;
 import org.openmetadata.catalog.security.Authorizer;
 import org.openmetadata.catalog.security.SecurityUtil;
 import org.openmetadata.catalog.type.EntityHistory;
-import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.Include;
 import org.openmetadata.catalog.util.EntityUtil;
 import org.openmetadata.catalog.util.RestUtil;
@@ -315,18 +311,8 @@ public class StorageServiceResource {
   public Response update(
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateStorageService update)
       throws IOException, ParseException {
-    EntityUtil.Fields fields = new EntityUtil.Fields(FIELD_LIST, FIELDS);
+    StorageService service = getService(update, securityContext);
-    EntityReference owner = null;
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(service));
-    StorageService service;
-    // Try to find the owner if entity exists
-    try {
-      service = dao.getByName(uriInfo, update.getName(), fields);
-      owner = helper(service).validateOwnerOrNull();
-    } catch (EntityNotFoundException e) {
-      // This is a create request if entity is not found. ignore exception
-    }
-    service = getService(update, securityContext);
-    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, owner);
     PutResponse<StorageService> response = dao.createOrUpdate(uriInfo, service, true);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/topics/TopicResource.java

@@ -354,8 +354,9 @@ public class TopicResource {
   public Response createOrUpdate(
       @Context UriInfo uriInfo, @Context SecurityContext securityContext, @Valid CreateTopic create)
       throws IOException, ParseException {
-
     Topic topic = getTopic(securityContext, create);
+    ;
+    SecurityUtil.checkAdminRoleOrPermissions(authorizer, securityContext, dao.getOriginalOwner(topic));
     PutResponse<Topic> response = dao.createOrUpdate(uriInfo, topic);
     addHref(uriInfo, response.getEntity());
     return response.toResponse();

catalog-rest-service/src/main/resources/json/schema/api/data/createTable.json

@@ -38,7 +38,7 @@
     },
     "database": {
       "description": "Database corresponding to this table",
-      "$ref": "../../type/basic.json#/definitions/uuid",
+      "$ref": "../../type/entityReference.json",
       "default": null
     },
     "tags": {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/EntityResourceTest.java

@@ -99,6 +99,7 @@ import org.openmetadata.catalog.entity.teams.User;
 import org.openmetadata.catalog.exception.CatalogExceptionMessage;
 import org.openmetadata.catalog.jdbi3.ChartRepository.ChartEntityInterface;
 import org.openmetadata.catalog.jdbi3.DashboardServiceRepository.DashboardServiceEntityInterface;
+import org.openmetadata.catalog.jdbi3.DatabaseRepository.DatabaseEntityInterface;
 import org.openmetadata.catalog.jdbi3.DatabaseServiceRepository.DatabaseServiceEntityInterface;
 import org.openmetadata.catalog.jdbi3.MessagingServiceRepository.MessagingServiceEntityInterface;
 import org.openmetadata.catalog.jdbi3.PipelineServiceRepository.PipelineServiceEntityInterface;
@@ -191,6 +192,7 @@ public abstract class EntityResourceTest<T, K> extends CatalogApplicationTest {
   public static List<EntityReference> CHART_REFERENCES;
 
   public static Database DATABASE;
+  public static EntityReference DATABASE_REFERENCE;
 
   public static List<Column> COLUMNS;
 
@@ -362,6 +364,7 @@ public abstract class EntityResourceTest<T, K> extends CatalogApplicationTest {
     DatabaseResourceTest databaseResourceTest = new DatabaseResourceTest();
     CreateDatabase create = databaseResourceTest.createRequest(test).withService(SNOWFLAKE_REFERENCE);
     DATABASE = databaseResourceTest.createAndCheckEntity(create, ADMIN_AUTH_HEADERS);
+    DATABASE_REFERENCE = new DatabaseEntityInterface(DATABASE).getEntityReference();
 
     COLUMNS =
         Arrays.asList(
@@ -818,6 +821,27 @@ public abstract class EntityResourceTest<T, K> extends CatalogApplicationTest {
     checkOwnerOwns(USER_OWNER1, entityInterface.getId(), true);
   }
 
+  @Test
+  void put_entityUpdate_as_non_owner_4xx(TestInfo test) throws IOException {
+    if (!supportsOwner) {
+      return; // Entity doesn't support ownership
+    }
+
+    // Create an entity with owner
+    K request = createRequest(getEntityName(test), "description", "displayName", USER_OWNER1);
+    createAndCheckEntity(request, ADMIN_AUTH_HEADERS);
+
+    // Update description and remove owner as non-owner
+    // Expect to throw an exception since only owner or admin can update resource
+    K updateRequest = createRequest(getEntityName(test), "newdescription", "displayName", null);
+    HttpResponseException exception =
+        assertThrows(
+            HttpResponseException.class,
+            () -> updateAndCheckEntity(updateRequest, OK, TEST_AUTH_HEADERS, UpdateType.NO_CHANGE, null));
+    TestUtils.assertResponse(
+        exception, FORBIDDEN, "Principal: CatalogPrincipal{name='test'} " + "does not have permissions");
+  }
+
   @Test
   void put_entityNullDescriptionUpdate_200(TestInfo test) throws IOException {
     // Create entity with null description

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/databases/TableResourceTest.java

@@ -301,7 +301,8 @@ public class TableResourceTest extends EntityResourceTest<Table, CreateTable> {
 
   @Test
   void post_tableWithInvalidDatabase_404(TestInfo test) {
-    CreateTable create = createRequest(test).withDatabase(NON_EXISTENT_ENTITY);
+    EntityReference database = new EntityReference().withId(NON_EXISTENT_ENTITY).withType(Entity.DATABASE);
+    CreateTable create = createRequest(test).withDatabase(database);
     HttpResponseException exception =
         assertThrows(HttpResponseException.class, () -> createEntity(create, ADMIN_AUTH_HEADERS));
     assertResponse(exception, NOT_FOUND, CatalogExceptionMessage.entityNotFound(Entity.DATABASE, NON_EXISTENT_ENTITY));
@@ -1400,7 +1401,7 @@ public class TableResourceTest extends EntityResourceTest<Table, CreateTable> {
         new TableConstraint().withConstraintType(ConstraintType.UNIQUE).withColumns(List.of(COLUMNS.get(0).getName()));
     return new CreateTable()
         .withName(name)
-        .withDatabase(DATABASE.getId())
+        .withDatabase(DATABASE_REFERENCE)
         .withColumns(COLUMNS)
         .withTableConstraints(List.of(constraint))
         .withDescription(description)
@@ -1468,15 +1469,15 @@ public class TableResourceTest extends EntityResourceTest<Table, CreateTable> {
     // Entity specific validation
     assertEquals(expected.getTableType(), patched.getTableType());
     assertColumns(expected.getColumns(), patched.getColumns());
-    validateDatabase(expected.getDatabase().getId(), patched.getDatabase());
+    validateDatabase(expected.getDatabase(), patched.getDatabase());
     assertEquals(expected.getTableConstraints(), patched.getTableConstraints());
     TestUtils.validateTags(expected.getTags(), patched.getTags());
     TestUtils.validateEntityReference(expected.getFollowers());
   }
 
-  private void validateDatabase(UUID expectedDatabaseId, EntityReference database) {
+  private void validateDatabase(EntityReference expectedDatabase, EntityReference database) {
     TestUtils.validateEntityReference(database);
-    assertEquals(expectedDatabaseId, database.getId());
+    assertEquals(expectedDatabase.getId(), database.getId());
   }
 
   @Override

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/mlmodels/MlModelResourceTest.java

@@ -116,7 +116,8 @@ public class MlModelResourceTest extends EntityResourceTest<MlModel, CreateMlMod
             dashboardResourceTest.createRequest(test).withCharts(null), ADMIN_AUTH_HEADERS);
     DASHBOARD_REFERENCE = new DashboardEntityInterface(DASHBOARD).getEntityReference();
 
-    CreateTable createTable = new CreateTable().withName("myTable").withDatabase(DATABASE.getId()).withColumns(COLUMNS);
+    CreateTable createTable =
+        new CreateTable().withName("myTable").withDatabase(DATABASE_REFERENCE).withColumns(COLUMNS);
 
     TableResourceTest tableResourceTest = new TableResourceTest();
     TABLE = tableResourceTest.createAndCheckEntity(createTable, ADMIN_AUTH_HEADERS);

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/services/DashboardServiceResourceTest.java

@@ -14,12 +14,10 @@
 package org.openmetadata.catalog.resources.services;
 
 import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
-import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static javax.ws.rs.core.Response.Status.OK;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.openmetadata.catalog.util.TestUtils.ADMIN_AUTH_HEADERS;
-import static org.openmetadata.catalog.util.TestUtils.TEST_AUTH_HEADERS;
 import static org.openmetadata.catalog.util.TestUtils.getPrincipal;
 
 import java.io.IOException;
@@ -191,21 +189,6 @@ public class DashboardServiceResourceTest extends EntityResourceTest<DashboardSe
     updateAndCheckEntity(update, OK, ADMIN_AUTH_HEADERS, UpdateType.MINOR_UPDATE, change);
   }
 
-  @Test
-  void put_update_as_non_owner_401(TestInfo test) throws IOException {
-    createAndCheckEntity(
-        createRequest(test).withDescription(null).withIngestionSchedule(null).withOwner(USER_OWNER1),
-        ADMIN_AUTH_HEADERS);
-
-    // Update dashboard description and ingestion service that are null
-    HttpResponseException exception =
-        assertThrows(
-            HttpResponseException.class,
-            () -> updateAndCheckEntity(createRequest(test), OK, TEST_AUTH_HEADERS, UpdateType.NO_CHANGE, null));
-    TestUtils.assertResponse(
-        exception, FORBIDDEN, "Principal: CatalogPrincipal{name='test'} " + "does not have permissions");
-  }
-
   @Override
   public CreateDashboardService createRequest(
       String name, String description, String displayName, EntityReference owner) {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/services/DatabaseServiceResourceTest.java

@@ -14,13 +14,11 @@
 package org.openmetadata.catalog.resources.services;
 
 import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
-import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static javax.ws.rs.core.Response.Status.OK;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.openmetadata.catalog.Entity.helper;
 import static org.openmetadata.catalog.util.TestUtils.ADMIN_AUTH_HEADERS;
-import static org.openmetadata.catalog.util.TestUtils.TEST_AUTH_HEADERS;
 import static org.openmetadata.catalog.util.TestUtils.getPrincipal;
 
 import java.io.IOException;
@@ -174,19 +172,6 @@ public class DatabaseServiceResourceTest extends EntityResourceTest<DatabaseServ
     assertEquals(airflowPipeline.getFullyQualifiedName(), expectedPipeline.getName());
   }
 
-  @Test
-  void put_update_as_non_owner_401(TestInfo test) throws IOException {
-    createAndCheckEntity(createRequest(test).withDescription(null).withOwner(USER_OWNER1), ADMIN_AUTH_HEADERS);
-
-    // Update as non owner should be forbidden
-    HttpResponseException exception =
-        assertThrows(
-            HttpResponseException.class,
-            () -> updateAndCheckEntity(createRequest(test), OK, TEST_AUTH_HEADERS, UpdateType.MINOR_UPDATE, null));
-    TestUtils.assertResponse(
-        exception, FORBIDDEN, "Principal: CatalogPrincipal{name='test'} " + "does not have permissions");
-  }
-
   @Override
   public CreateDatabaseService createRequest(
       String name, String description, String displayName, EntityReference owner) {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/services/MessagingServiceResourceTest.java

@@ -14,13 +14,11 @@
 package org.openmetadata.catalog.resources.services;
 
 import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
-import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static javax.ws.rs.core.Response.Status.OK;
 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.openmetadata.catalog.util.TestUtils.ADMIN_AUTH_HEADERS;
-import static org.openmetadata.catalog.util.TestUtils.TEST_AUTH_HEADERS;
 
 import java.io.IOException;
 import java.net.URI;
@@ -217,21 +215,6 @@ public class MessagingServiceResourceTest extends EntityResourceTest<MessagingSe
     updateAndCheckEntity(update, OK, ADMIN_AUTH_HEADERS, UpdateType.MINOR_UPDATE, change);
   }
 
-  @Test
-  void put_update_as_non_owner_401(TestInfo test) throws IOException {
-    createAndCheckEntity(
-        createRequest(test).withDescription(null).withIngestionSchedule(null).withOwner(USER_OWNER1),
-        ADMIN_AUTH_HEADERS);
-
-    // Update messaging description as non owner and expect exception
-    HttpResponseException exception =
-        assertThrows(
-            HttpResponseException.class,
-            () -> updateAndCheckEntity(createRequest(test), OK, TEST_AUTH_HEADERS, UpdateType.NO_CHANGE, null));
-    TestUtils.assertResponse(
-        exception, FORBIDDEN, "Principal: CatalogPrincipal{name='test'} " + "does not have permissions");
-  }
-
   @Override
   public CreateMessagingService createRequest(
       String name, String description, String displayName, EntityReference owner) {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/services/PipelineServiceResourceTest.java

@@ -14,12 +14,10 @@
 package org.openmetadata.catalog.resources.services;
 
 import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
-import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static javax.ws.rs.core.Response.Status.OK;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.openmetadata.catalog.util.TestUtils.ADMIN_AUTH_HEADERS;
-import static org.openmetadata.catalog.util.TestUtils.TEST_AUTH_HEADERS;
 import static org.openmetadata.catalog.util.TestUtils.getPrincipal;
 
 import java.io.IOException;
@@ -197,21 +195,6 @@ public class PipelineServiceResourceTest extends EntityResourceTest<PipelineServ
     updateAndCheckEntity(update, OK, ADMIN_AUTH_HEADERS, UpdateType.MINOR_UPDATE, change);
   }
 
-  @Test
-  void put_update_as_non_owner_401(TestInfo test) throws IOException {
-    createAndCheckEntity(
-        createRequest(test).withDescription(null).withIngestionSchedule(null).withOwner(USER_OWNER1),
-        ADMIN_AUTH_HEADERS);
-
-    // Update pipeline description and ingestion service that are null
-    HttpResponseException exception =
-        assertThrows(
-            HttpResponseException.class,
-            () -> updateAndCheckEntity(createRequest(test), OK, TEST_AUTH_HEADERS, UpdateType.NO_CHANGE, null));
-    TestUtils.assertResponse(
-        exception, FORBIDDEN, "Principal: CatalogPrincipal{name='test'} " + "does not have permissions");
-  }
-
   @Override
   public CreatePipelineService createRequest(
       String name, String description, String displayName, EntityReference owner) {

catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/services/StorageServiceResourceTest.java

@@ -13,12 +13,8 @@
 
 package org.openmetadata.catalog.resources.services;
 
-import static javax.ws.rs.core.Response.Status.FORBIDDEN;
-import static javax.ws.rs.core.Response.Status.OK;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.openmetadata.catalog.util.TestUtils.ADMIN_AUTH_HEADERS;
-import static org.openmetadata.catalog.util.TestUtils.TEST_AUTH_HEADERS;
 import static org.openmetadata.catalog.util.TestUtils.getPrincipal;
 
 import java.io.IOException;
@@ -37,7 +33,6 @@ import org.openmetadata.catalog.type.EntityReference;
 import org.openmetadata.catalog.type.StorageServiceType;
 import org.openmetadata.catalog.util.EntityInterface;
 import org.openmetadata.catalog.util.TestUtils;
-import org.openmetadata.catalog.util.TestUtils.UpdateType;
 
 @Slf4j
 public class StorageServiceResourceTest extends EntityResourceTest<StorageService, CreateStorageService> {
@@ -70,19 +65,6 @@ public class StorageServiceResourceTest extends EntityResourceTest<StorageServic
     // TODO add more tests for different fields
   }
 
-  @Test
-  void put_update_as_non_owner_401(TestInfo test) throws IOException {
-    createAndCheckEntity(createRequest(test).withDescription(null).withOwner(USER_OWNER1), ADMIN_AUTH_HEADERS);
-
-    // Update storage description and ingestion service as non owner should be forbidden
-    HttpResponseException exception =
-        assertThrows(
-            HttpResponseException.class,
-            () -> updateAndCheckEntity(createRequest(test), OK, TEST_AUTH_HEADERS, UpdateType.NO_CHANGE, null));
-    TestUtils.assertResponse(
-        exception, FORBIDDEN, "Principal: CatalogPrincipal{name='test'} " + "does not have permissions");
-  }
-
   @Override
   public CreateStorageService createRequest(
       String name, String description, String displayName, EntityReference owner) {

ingestion/src/metadata/ingestion/sink/metadata_rest.py

@@ -137,6 +137,7 @@ class MetadataRestSink(Sink[Entity]):
                 ),
             )
             db = self.metadata.create_or_update(db_request)
+            db_ref = EntityReference(id=db.id, name=db.name.__root__, type="database")
             if db_and_table.table.description is not None:
                 db_and_table.table.description = db_and_table.table.description.strip()
 
@@ -145,7 +146,7 @@ class MetadataRestSink(Sink[Entity]):
                 tableType=db_and_table.table.tableType,
                 columns=db_and_table.table.columns,
                 description=db_and_table.table.description,
-                database=db.id,
+                database=db_ref,
             )
             if db_and_table.table.viewDefinition:
                 table_request.viewDefinition = (

ingestion/tests/integration/lineage/airflow/test_airflow_lineage.py

@@ -83,9 +83,11 @@ class AirflowLineageTest(TestCase):
 
         cls.create_db_entity = cls.metadata.create_or_update(data=cls.create_db)
 
+        cls.db_reference = EntityReference(id=cls.create_db_entity.id, name="test-db", type="database")
+
         cls.create = CreateTableRequest(
             name="lineage-test",
-            database=cls.create_db_entity.id,
+            database=cls.db_reference,
             columns=[Column(name="id", dataType=DataType.BIGINT)],
         )
 

ingestion/tests/integration/ometa/test_ometa_lineage_api.py

@@ -87,9 +87,11 @@ class OMetaLineageTest(TestCase):
 
         cls.create_db_entity = cls.metadata.create_or_update(data=cls.create_db)
 
+        cls.db_reference = EntityReference(id=cls.create_db_entity.id, name="test-db", type="database")
+
         cls.table = CreateTableRequest(
             name="test",
-            database=cls.create_db_entity.id,
+            database=cls.db_reference,
             columns=[Column(name="id", dataType=DataType.BIGINT)],
         )
 

ingestion/tests/integration/ometa/test_ometa_model_api.py

@@ -203,16 +203,18 @@ class OMetaModelTest(TestCase):
         )
         create_db_entity = self.metadata.create_or_update(data=create_db)
 
+        db_reference = EntityReference(id=create_db_entity.id, name="test-db-ml", type="database")
+
         create_table1 = CreateTableRequest(
             name="test-ml",
-            database=create_db_entity.id,
+            database=db_reference,
             columns=[Column(name="education", dataType=DataType.STRING)],
         )
         table1_entity = self.metadata.create_or_update(data=create_table1)
 
         create_table2 = CreateTableRequest(
             name="another_test-ml",
-            database=create_db_entity.id,
+            database=db_reference,
             columns=[Column(name="age", dataType=DataType.INT)],
         )
         table2_entity = self.metadata.create_or_update(data=create_table2)

ingestion/tests/integration/ometa/test_ometa_table_api.py

@@ -85,17 +85,19 @@ class OMetaTableTest(TestCase):
 
         cls.create_db_entity = cls.metadata.create_or_update(data=cls.create_db)
 
+        cls.db_reference = EntityReference(id=cls.create_db_entity.id, name="test-db", type="database")
+
         cls.entity = Table(
             id=uuid.uuid4(),
             name="test",
-            database=EntityReference(id=cls.create_db_entity.id, type="database"),
+            database=cls.db_reference,
             fullyQualifiedName="test-service-table.test-db.test",
             columns=[Column(name="id", dataType=DataType.BIGINT)],
         )
 
         cls.create = CreateTableRequest(
             name="test",
-            database=cls.create_db_entity.id,
+            database=cls.db_reference,
             columns=[Column(name="id", dataType=DataType.BIGINT)],
         )
 
@@ -151,7 +153,7 @@ class OMetaTableTest(TestCase):
         res = self.metadata.create_or_update(data=updated_entity)
 
         # Same ID, updated algorithm
-        self.assertEqual(res.database.id, updated_entity.database)
+        self.assertEqual(res.database.id, updated_entity.database.id)
         self.assertEqual(res_create.id, res.id)
         self.assertEqual(res.owner.id, self.user.id)
 
@@ -310,7 +312,7 @@ class OMetaTableTest(TestCase):
 
         another_table = CreateTableRequest(
             name="another-test",
-            database=self.create_db_entity.id,
+            database=self.db_reference,
             columns=[Column(name="another_id", dataType=DataType.BIGINT)],
         )
         another_res = self.metadata.create_or_update(another_table)

ingestion/tests/integration/source/hive/test_hive_crud.py

@@ -76,7 +76,8 @@ def create_delete_table(client: OpenMetadata, databases: List[Database]):
         Column(name="id", dataType="INT", dataLength=1),
         Column(name="name", dataType="VARCHAR", dataLength=1),
     ]
-    table = CreateTableRequest(name="test1", columns=columns, database=databases[0].id)
+    db_ref = EntityReference(id=databases[0].id, name=databases[0].name.__root__, type="database")
+    table = CreateTableRequest(name="test1", columns=columns, database=db_ref)
     created_table = client.create_or_update(table)
     if table.name.__root__ == created_table.name.__root__:
         client.delete(entity=Table, entity_id=str(created_table.id.__root__))

ingestion/tests/integration/source/mssql/test_mssql_crud.py

@@ -64,7 +64,8 @@ def create_delete_table(client):
         Column(name="id", columnDataType="INT"),
         Column(name="name", columnDataType="VARCHAR"),
     ]
-    table = CreateTableRequest(name="test1", columns=columns, database=databases[0].id)
+    db_ref = EntityReference(id=databases[0].id, name=databases[0].name.__root__, type="database")
+    table = CreateTableRequest(name="test1", columns=columns, database=db_ref)
     created_table = client.create_or_update_table(table)
     if table.name.__root__ == created_table.name.__root__:
         requests.delete(

ingestion/tests/integration/source/mysql/test_mysql_crud.py

@@ -45,7 +45,8 @@ def create_delete_table(client: OpenMetadata):
         Column(name="id", dataType="INT", dataLength=1),
         Column(name="name", dataType="VARCHAR", dataLength=1),
     ]
-    table = CreateTableRequest(name="test1", columns=columns, database=databases[0].id)
+    db_ref = EntityReference(id=databases[0].id, name=databases[0].name.__root__, type="database")
+    table = CreateTableRequest(name="test1", columns=columns, database=db_ref)
     created_table = client.create_or_update(table)
     if table.name.__root__ == created_table.name.__root__:
         requests.delete(

ingestion/tests/integration/source/postgres/test_postgres_crud.py

@@ -87,8 +87,9 @@ def test_create_table_service(catalog_service):
         service=EntityReference(id=postgres_dbservice.id, type="databaseService"),
     )
     created_database = client.create_database(create_database_request)
+    db_ref = EntityReference(id=created_database.id.__root__, name=created_database.name.__root__, type="database")
     table = CreateTableRequest(
-        name=table_name, columns=columns, database=created_database.id.__root__
+        name=table_name, columns=columns, database=db_ref
     )
     created_table = client.create_or_update_table(table)
     if created_database and created_table:

ingestion/tests/integration/source/trino/test_trino_crud.py

@@ -81,8 +81,9 @@ def create_delete_table(client: OpenMetadata, databases: List[Database]):
         Column(name="name", dataType="VARCHAR", dataLength=1),
     ]
     print(databases[0])
+    db_ref = EntityReference(id=databases[0].id.__root__, name=databases[0].name.__root__, type="database")
     table = CreateTableRequest(
-        name="test1", columns=columns, database=databases[0].id.__root__
+        name="test1", columns=columns, database=db_ref
     )
     created_table = client.create_or_update(table)
     if table.name.__root__ == created_table.name.__root__: