yujunjun/OpenMetadata
1
0
Fork 0
mirror of https://github.com/open-metadata/OpenMetadata.git synced 2026-01-08 13:36:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add queryUsage and other fields for bulkQuery Fetching (#23087)

Browse Source 
* Add queryUsage and other fields for bulkQuery Fetching

* remove additional comments
This commit is contained in:
sonika-shah 2025-08-26 10:10:53 +05:30 committed by GitHub
parent f739ba2355
commit 2ddc61a0ca
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 162 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/QueryRepository.java
View File

@ -10,7 +10,9 @@ import jakarta.ws.rs.core.Response;
import jakarta.ws.rs.core.UriInfo;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
@ -84,6 +86,10 @@ public class QueryRepository extends EntityRepository<Query> {
// Bulk fetch and set services for all queries first
fetchAndSetServices(entities);
// Bulk fetch and set query-specific fields
fetchAndSetQueryUsage(entities, fields);
fetchAndSetQueryUsers(entities, fields);
// Then call parent's implementation which handles standard fields
super.setFieldsInBulk(fields, entities);
}
@ -117,6 +123,62 @@ public class QueryRepository extends EntityRepository<Query> {
}
}
private void fetchAndSetQueryUsage(List<Query> queries, EntityUtil.Fields fields) {
if (!fields.contains(QUERY_USED_IN_FIELD) || queries == null || queries.isEmpty()) {
return;
}
List<String> queryIds = queries.stream().map(q -> q.getId().toString()).toList();
List<CollectionDAO.EntityRelationshipObject> relationships =
daoCollection
.relationshipDAO()
.findFromBatch(queryIds, Entity.QUERY, Relationship.MENTIONED_IN.ordinal());
// Group relationships by query ID
Map<UUID, List<EntityReference>> queryUsageMap = new HashMap<>();
for (CollectionDAO.EntityRelationshipObject record : relationships) {
UUID queryId = UUID.fromString(record.getToId());
EntityReference entityRef =
Entity.getEntityReferenceById(
record.getFromEntity(), UUID.fromString(record.getFromId()), Include.ALL);
queryUsageMap.computeIfAbsent(queryId, k -> new ArrayList<>()).add(entityRef);
}
queries.forEach(
query -> {
List<EntityReference> usage = queryUsageMap.getOrDefault(query.getId(), List.of());
query.setQueryUsedIn(usage);
});
}
private void fetchAndSetQueryUsers(List<Query> queries, EntityUtil.Fields fields) {
if (!fields.contains("users") || queries == null || queries.isEmpty()) {
return;
}
List<String> queryIds = queries.stream().map(q -> q.getId().toString()).toList();
List<CollectionDAO.EntityRelationshipObject> relationships =
daoCollection
.relationshipDAO()
.findFromBatch(queryIds, Entity.QUERY, Relationship.USES.ordinal());
// Group relationships by query ID
Map<UUID, List<EntityReference>> queryUsersMap = new HashMap<>();
for (CollectionDAO.EntityRelationshipObject record : relationships) {
UUID queryId = UUID.fromString(record.getToId());
EntityReference entityRef =
Entity.getEntityReferenceById(
record.getFromEntity(), UUID.fromString(record.getFromId()), Include.ALL);
queryUsersMap.computeIfAbsent(queryId, k -> new ArrayList<>()).add(entityRef);
}
queries.forEach(
query -> {
List<EntityReference> users = queryUsersMap.getOrDefault(query.getId(), List.of());
query.withUsers(users);
});
}
public List<EntityReference> getQueryUsage(Query queryEntity) {
return queryEntity == null
? Collections.emptyList()

100
openmetadata-service/src/test/java/org/openmetadata/service/resources/query/QueryResourceTest.java
View File

@ -18,6 +18,7 @@ import jakarta.ws.rs.core.Response;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.RandomStringUtils;
@ -331,6 +332,105 @@ public class QueryResourceTest extends EntityResourceTest<Query, CreateQuery> {
assertEquals(updatedQuery.getChecksum(), EntityUtil.hash(updatedQuery.getQuery()));
}
@Test
void test_batchFetchQueryFields(TestInfo test) throws IOException {
// Test bulk fetching of query relationships using setFieldsInBulk method
String testName = getEntityName(test);
// Create additional table for testing queryUsedIn relationships
TableResourceTest tableResourceTest = new TableResourceTest();
List<Column> columns = List.of(TableResourceTest.getColumn(C1, ColumnDataType.INT, null));
CreateTable tableCreate =
tableResourceTest.createRequest(test).withName(testName + "_table").withColumns(columns);
Table testTable = tableResourceTest.createAndCheckEntity(tableCreate, ADMIN_AUTH_HEADERS);
// Create queries with different relationship patterns
CreateQuery query1Create =
createRequest(testName + "_query1")
.withQueryUsedIn(List.of(TABLE_REF, testTable.getEntityReference()))
.withUsers(List.of(USER1.getName(), USER2.getName()));
Query query1 = createAndCheckEntity(query1Create, ADMIN_AUTH_HEADERS);
CreateQuery query2Create =
createRequest(testName + "_query2")
.withQueryUsedIn(List.of(testTable.getEntityReference()))
.withUsers(List.of(USER1.getName()));
Query query2 = createAndCheckEntity(query2Create, ADMIN_AUTH_HEADERS);
CreateQuery query3Create =
createRequest(testName + "_query3")
.withQueryUsedIn(List.of())
.withUsers(List.of(USER2.getName()));
Query query3 = createAndCheckEntity(query3Create, ADMIN_AUTH_HEADERS);
try {
// Test 1: Bulk fetch with all fields - should populate all relationships
ResultList<Query> allFieldsResult = getQueries(100, "*", true, ADMIN_AUTH_HEADERS);
Query fetchedQuery1 = findQueryInResults(allFieldsResult, query1.getId());
Query fetchedQuery2 = findQueryInResults(allFieldsResult, query2.getId());
Query fetchedQuery3 = findQueryInResults(allFieldsResult, query3.getId());
// Verify bulk fetchers populated queryUsedIn correctly
assertNotNull(fetchedQuery1, "Query1 should be found");
assertListNotNull(fetchedQuery1.getQueryUsedIn());
assertEquals(2, fetchedQuery1.getQueryUsedIn().size(), "Query1 should have 2 queryUsedIn");
assertListNotNull(fetchedQuery1.getUsers());
assertEquals(2, fetchedQuery1.getUsers().size(), "Query1 should have 2 users");
assertNotNull(fetchedQuery2, "Query2 should be found");
assertListNotNull(fetchedQuery2.getQueryUsedIn());
assertEquals(1, fetchedQuery2.getQueryUsedIn().size(), "Query2 should have 1 queryUsedIn");
assertListNotNull(fetchedQuery2.getUsers());
assertEquals(1, fetchedQuery2.getUsers().size(), "Query2 should have 1 user");
assertNotNull(fetchedQuery3, "Query3 should be found");
assertListNotNull(fetchedQuery3.getQueryUsedIn());
assertEquals(0, fetchedQuery3.getQueryUsedIn().size(), "Query3 should have 0 queryUsedIn");
assertListNotNull(fetchedQuery3.getUsers());
assertEquals(1, fetchedQuery3.getUsers().size(), "Query3 should have 1 user");
// Test 2: Fetch only queryUsedIn field - should only populate queryUsedIn
ResultList<Query> queryUsedInOnly = getQueries(100, "queryUsedIn", true, ADMIN_AUTH_HEADERS);
Query queryUsedInResult = findQueryInResults(queryUsedInOnly, query1.getId());
assertNotNull(queryUsedInResult, "Query should be found with queryUsedIn field");
assertListNotNull(queryUsedInResult.getQueryUsedIn());
assertEquals(2, queryUsedInResult.getQueryUsedIn().size());
assertListNull(queryUsedInResult.getUsers()); // Should be null - not requested
// Test 3: Fetch only users field - should only populate users
ResultList<Query> usersOnly = getQueries(100, "users", true, ADMIN_AUTH_HEADERS);
Query usersResult = findQueryInResults(usersOnly, query1.getId());
assertNotNull(usersResult, "Query should be found with users field");
assertListNotNull(usersResult.getUsers());
assertEquals(2, usersResult.getUsers().size());
assertListNull(usersResult.getQueryUsedIn()); // Should be null - not requested
// Test 4: Fetch without relationship fields - should not populate relationships
ResultList<Query> noRelFields = getQueries(100, "name,query", true, ADMIN_AUTH_HEADERS);
Query noRelResult = findQueryInResults(noRelFields, query1.getId());
assertNotNull(noRelResult, "Query should be found without relationship fields");
assertListNull(noRelResult.getQueryUsedIn());
assertListNull(noRelResult.getUsers());
} finally {
// Cleanup test queries
deleteEntity(query1.getId(), ADMIN_AUTH_HEADERS);
deleteEntity(query2.getId(), ADMIN_AUTH_HEADERS);
deleteEntity(query3.getId(), ADMIN_AUTH_HEADERS);
tableResourceTest.deleteEntity(testTable.getId(), ADMIN_AUTH_HEADERS);
}
}
private Query findQueryInResults(ResultList<Query> results, UUID queryId) {
return results.getData().stream()
.filter(q -> q.getId().equals(queryId))
.findFirst()
.orElse(null);
}
public ResultList<Query> getQueries(
Integer limit, String fields, Boolean includeAll, Map<String, String> authHeaders)
throws HttpResponseException {