diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DatabaseRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DatabaseRepository.java index 3a993751b2d..8137700e1ae 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DatabaseRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DatabaseRepository.java @@ -329,10 +329,17 @@ public class DatabaseRepository extends EntityRepository { return; } - EntityReference service = getContainer(entities.get(0).getId()); - for (Database database : entities) { - database.setService(service); - } + // Use batch fetch to get correct service for each database + var serviceMap = batchFetchServices(entities); + + // Set the correct service for each database + entities.forEach( + database -> { + EntityReference service = serviceMap.get(database.getId()); + if (service != null) { + database.setService(service); + } + }); } private Map> batchFetchDatabaseSchemas(List databases) { diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseResourceTest.java index 2efc28be8bc..b375f5a8514 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseResourceTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/databases/DatabaseResourceTest.java @@ -44,8 +44,10 @@ import java.util.Map; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; import org.apache.http.client.HttpResponseException; +import org.junit.jupiter.api.MethodOrderer; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.TestInfo; +import org.junit.jupiter.api.TestMethodOrder; import org.openmetadata.csv.CsvUtil; import org.openmetadata.csv.EntityCsv; import org.openmetadata.schema.api.data.CreateDatabase; @@ -68,6 +70,7 @@ import org.openmetadata.service.util.ResultList; import org.openmetadata.service.util.TestUtils; @Slf4j +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) public class DatabaseResourceTest extends EntityResourceTest { public DatabaseResourceTest() { super( @@ -428,4 +431,87 @@ public class DatabaseResourceTest extends EntityResourceTest existingDatabases = + listEntities(Map.of("fields", "service", "limit", "50"), ADMIN_AUTH_HEADERS); + + // Group by service to check if each database has its correct service + Map> dbByService = new HashMap<>(); + for (Database db : existingDatabases.getData()) { + assertNotNull(db.getService(), "Database " + db.getName() + " should have a service"); + String serviceName = db.getService().getName(); + dbByService.computeIfAbsent(serviceName, k -> new ArrayList<>()).add(db); + } + + // If we have databases from different services, verify each has the correct service + if (dbByService.size() > 1) { + for (Map.Entry> entry : dbByService.entrySet()) { + String serviceName = entry.getKey(); + for (Database db : entry.getValue()) { + assertEquals( + serviceName, + db.getService().getName(), + "Database " + db.getName() + " should have service " + serviceName); + } + } + LOG.info( + "Verified databases maintain correct services across {} different services", + dbByService.size()); + } + + // Now test with new databases in the default container to ensure entity count isn't affected + String timestamp = String.valueOf(System.currentTimeMillis()); + CreateDatabase createDb1 = createRequest("bulkTestDb1_" + timestamp); + CreateDatabase createDb2 = createRequest("bulkTestDb2_" + timestamp); + + Database db1 = createAndCheckEntity(createDb1, ADMIN_AUTH_HEADERS); + Database db2 = createAndCheckEntity(createDb2, ADMIN_AUTH_HEADERS); + + try { + // Fetch with service field + ResultList databases = + listEntities( + Map.of("fields", "service", "service", getContainer().getFullyQualifiedName()), + ADMIN_AUTH_HEADERS); + + // Find our databases + Database foundDb1 = + databases.getData().stream() + .filter(db -> db.getId().equals(db1.getId())) + .findFirst() + .orElse(null); + Database foundDb2 = + databases.getData().stream() + .filter(db -> db.getId().equals(db2.getId())) + .findFirst() + .orElse(null); + + assertNotNull(foundDb1, "Database 1 should be found"); + assertNotNull(foundDb2, "Database 2 should be found"); + + // Both should have the same service (SNOWFLAKE_REFERENCE) since they're in the same container + assertEquals(getContainer().getId(), foundDb1.getService().getId()); + assertEquals(getContainer().getId(), foundDb2.getService().getId()); + assertEquals(getContainer().getName(), foundDb1.getService().getName()); + assertEquals(getContainer().getName(), foundDb2.getService().getName()); + + } finally { + // Clean up - these will be properly handled by recursive delete since they're in the default + // container + deleteEntity(db1.getId(), ADMIN_AUTH_HEADERS); + deleteEntity(db2.getId(), ADMIN_AUTH_HEADERS); + } + } } 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 3e85fd0ac73..29936c8940d 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 @@ -43,8 +43,11 @@ import java.util.Map; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; import org.apache.http.client.HttpResponseException; +import org.junit.jupiter.api.MethodOrderer; +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.csv.CsvUtil; import org.openmetadata.csv.EntityCsv; import org.openmetadata.schema.api.data.CreateDatabase; @@ -75,6 +78,7 @@ import org.openmetadata.service.util.FullyQualifiedName; import org.openmetadata.service.util.ResultList; import org.openmetadata.service.util.TestUtils; +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) @Slf4j public class DatabaseSchemaResourceTest extends EntityResourceTest { @@ -753,4 +757,57 @@ public class DatabaseSchemaResourceTest public void assertFieldChange(String fieldName, Object expected, Object actual) { assertCommonFieldChange(fieldName, expected, actual); } + + @Order(1) + @Test + void testSchemaServiceInheritanceFromDatabase(TestInfo test) throws IOException { + // This test verifies that schemas correctly inherit service from their database + // when fetched in bulk with the service field + + // Use the existing DATABASE which already has a service set + Database database = DATABASE; + assertNotNull(database.getService(), "Test database should have a service"); + + // Create a schema in the database + String uniqueName = "serviceInheritanceTest" + System.currentTimeMillis(); + CreateDatabaseSchema createSchema = + createRequest(uniqueName).withDatabase(database.getFullyQualifiedName()); + DatabaseSchema schema = createAndCheckEntity(createSchema, ADMIN_AUTH_HEADERS); + + // Fetch schemas with the service field included + ResultList schemas = + listEntities( + Map.of( + "limit", + "10", + "fields", + "database,service", + "database", + database.getFullyQualifiedName()), + ADMIN_AUTH_HEADERS); + + assertNotNull(schemas.getData()); + assertTrue(schemas.getData().size() >= 1); + + // Find our created schema + DatabaseSchema foundSchema = + schemas.getData().stream() + .filter(s -> s.getId().equals(schema.getId())) + .findFirst() + .orElse(null); + + assertNotNull(foundSchema, "Created schema should be in the results"); + + // Verify the schema has the correct database and service + assertNotNull(foundSchema.getDatabase(), "Database should not be null"); + assertEquals(database.getId(), foundSchema.getDatabase().getId()); + + assertNotNull( + foundSchema.getService(), "Service should not be null - should be inherited from database"); + assertEquals(database.getService().getId(), foundSchema.getService().getId()); + assertEquals(database.getService().getName(), foundSchema.getService().getName()); + + // Clean up + deleteEntity(schema.getId(), ADMIN_AUTH_HEADERS); + } }