Add Data Quality Test Casae (#21764)

* Add Test Parameter Definition tool

* Add Dq create test case
This commit is contained in:
Mohit Yadav 2025-06-13 19:28:35 +05:30 committed by GitHub
parent 891ff4184d
commit c7e92d42f1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 254 additions and 2 deletions

View File

@ -0,0 +1,91 @@
package org.openmetadata.service.mcp.tools;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.openmetadata.schema.api.tests.CreateTestCase;
import org.openmetadata.schema.tests.TestCase;
import org.openmetadata.schema.tests.TestCaseParameterValue;
import org.openmetadata.service.Entity;
import org.openmetadata.service.jdbi3.TestCaseRepository;
import org.openmetadata.service.limits.Limits;
import org.openmetadata.service.resources.dqtests.TestCaseMapper;
import org.openmetadata.service.resources.feeds.MessageParser;
import org.openmetadata.service.security.Authorizer;
import org.openmetadata.service.security.auth.CatalogSecurityContext;
import org.openmetadata.service.util.JsonUtils;
import org.openmetadata.service.util.RestUtil;
public class CreateTestCaseTool implements McpTool {
private final TestCaseMapper testCaseMapper = new TestCaseMapper();
@Override
public Map<String, Object> execute(
Authorizer authorizer,
CatalogSecurityContext catalogSecurityContext,
Map<String, Object> params) {
try {
String testDefinitionName = (String) params.get("testDefinitionName");
String entityFqn = (String) params.get("entityFqn");
String entityType =
params.containsKey("entityType") ? (String) params.get("entityType") : "table";
String description =
params.containsKey("description")
? (String) params.get("description")
: "Test case created by MCP tool";
String name =
params.containsKey("name")
? (String) params.get("name")
: "TestCase_" + System.currentTimeMillis();
MessageParser.EntityLink entityLink = new MessageParser.EntityLink(entityType, entityFqn);
String entityLinkValue = entityLink.getLinkString();
List<TestCaseParameterValue> parameterValue =
params.containsKey("parameterValues")
? JsonUtils.readOrConvertValues(
params.get("parameterValues"), TestCaseParameterValue.class)
: new ArrayList<>();
TestCaseRepository repository =
(TestCaseRepository) Entity.getEntityRepository(Entity.TEST_CASE);
String updatedBy = catalogSecurityContext.getUserPrincipal().getName();
TestCase testCase =
getTestCase(
name, description, entityLinkValue, testDefinitionName, parameterValue, updatedBy);
repository.setFullyQualifiedName(testCase);
repository.prepare(testCase, false);
RestUtil.PutResponse<TestCase> response =
repository.createOrUpdate(null, testCase, updatedBy);
return JsonUtils.getMap(response);
} catch (Exception e) {
throw new RuntimeException("Failed to create test case: " + e.getMessage(), e);
}
}
private TestCase getTestCase(
String name,
String description,
String entityLinkValue,
String testDefinitionName,
List<TestCaseParameterValue> parameterValue,
String updatedBy) {
return testCaseMapper.createToEntity(
new CreateTestCase()
.withName(name)
.withDisplayName(name)
.withDescription(description)
.withEntityLink(entityLinkValue)
.withParameterValues(parameterValue)
.withComputePassedFailedRowCount(false)
.withUseDynamicAssertion(false)
.withTestDefinition(testDefinitionName),
updatedBy);
}
@Override
public Map<String, Object> execute(
Authorizer authorizer,
Limits limits,
CatalogSecurityContext catalogSecurityContext,
Map<String, Object> map) {
throw new UnsupportedOperationException("TestDefinition does not requires limit validation.");
}
}

View File

@ -57,6 +57,12 @@ public class DefaultToolContext {
case "patch_entity":
result = new PatchEntityTool().execute(authorizer, limits, securityContext, params);
break;
case "get_test_definitions":
result = new TestDefinitionsTool().execute(authorizer, securityContext, params);
break;
case "create_test_case":
result = new CreateTestCaseTool().execute(authorizer, securityContext, params);
break;
default:
return new McpSchema.CallToolResult(
List.of(
@ -80,7 +86,7 @@ public class DefaultToolContext {
403)))),
true);
} catch (Exception ex) {
LOG.error("Error executing tool: {}", ex.getMessage());
LOG.error("Error executing tool: ", ex);
return new McpSchema.CallToolResult(
List.of(
new McpSchema.TextContent(

View File

@ -0,0 +1,69 @@
package org.openmetadata.service.mcp.tools;
import java.util.Map;
import org.openmetadata.schema.tests.TestPlatform;
import org.openmetadata.schema.type.Include;
import org.openmetadata.schema.type.MetadataOperation;
import org.openmetadata.service.Entity;
import org.openmetadata.service.jdbi3.ListFilter;
import org.openmetadata.service.jdbi3.TestDefinitionRepository;
import org.openmetadata.service.limits.Limits;
import org.openmetadata.service.security.Authorizer;
import org.openmetadata.service.security.auth.CatalogSecurityContext;
import org.openmetadata.service.security.policyevaluator.OperationContext;
import org.openmetadata.service.security.policyevaluator.ResourceContext;
import org.openmetadata.service.util.JsonUtils;
public class TestDefinitionsTool implements McpTool {
@Override
public Map<String, Object> execute(
Authorizer authorizer,
CatalogSecurityContext catalogSecurityContext,
Map<String, Object> params) {
int limit = 10;
if (params.containsKey("limit")) {
Object limitObj = params.get("limit");
if (limitObj instanceof Number) {
limit = ((Number) limitObj).intValue();
} else if (limitObj instanceof String) {
limit = Integer.parseInt((String) limitObj);
}
}
String entityType =
params.containsKey("entityType") ? (String) params.get("entityType") : "TABLE";
String testPlatformParam =
params.containsKey("testPlatform")
? (String) params.get("testPlatform")
: TestPlatform.OPEN_METADATA.value();
String after = params.containsKey("after") ? (String) params.get("after") : null;
TestDefinitionRepository repository =
(TestDefinitionRepository) Entity.getEntityRepository(Entity.TEST_DEFINITION);
OperationContext listOperationContext =
new OperationContext(entityType, MetadataOperation.VIEW_BASIC);
authorizer.authorize(
catalogSecurityContext,
listOperationContext,
new ResourceContext<>(Entity.TEST_DEFINITION));
ListFilter filter = new ListFilter(Include.NON_DELETED);
if (entityType != null) {
filter.addQueryParam("entityType", entityType);
}
if (testPlatformParam != null) {
filter.addQueryParam("testPlatform", testPlatformParam);
}
return JsonUtils.getMap(
repository.listAfter(null, repository.getFields("*"), filter, limit, after));
}
@Override
public Map<String, Object> execute(
Authorizer authorizer,
Limits limits,
CatalogSecurityContext catalogSecurityContext,
Map<String, Object> map) {
throw new UnsupportedOperationException("TestDefinition does not requires limit validation.");
}
}

View File

@ -2,7 +2,7 @@
"tools": [
{
"name": "search_metadata",
"description": "Find your data and business terms in OpenMetadata. For example if the user asks to 'find tables that contain customers information', then 'customers' should be the query, and the entity_type should be 'table'. Here make sure to use 'Href' is available in result to create a hyperlink to the entity in OpenMetadata.",
"description": "Find your data and business terms in OpenMetadata. For example if the user asks to 'find tables that contain customers information', then 'customers' should be the query, and the entity_type should be 'table'.",
"parameters": {
"description": "The search query to find metadata in the OpenMetadata catalog, entity type could be table, topic etc. Limit can be used to paginate on the data.",
"type": "object",
@ -150,6 +150,92 @@
"patch"
]
}
},
{
"name": "get_test_definitions",
"description": "This tool can be used to get all the test definitions in the OpenMetadata. It returns a list of test definitions. These test definition can be used to create a test case for a table or a table's column. While creating column test for entity, column data type should be in supportedDataTypes of parameterDefinition.",
"parameters": {
"description": "These test definition are required to create a test case for table or a table's column.",
"type": "object",
"properties": {
"entityType": {
"type": "string",
"description": "Entity Type can be 'TABLE' for table asset/entity or 'COLUMN' for column level tests. Default is TABLE."
},
"testPlatform": {
"type": "string",
"description": "Default value can be 'OpenMetadata'. Other Platform can be 'OpenMetadata','GreatExpectations', 'DBT', 'Deequ', 'Soda', 'Other' if the user specifically gives the platform name."
},
"after": {
"type": "string",
"description": "This can be used to paginate the results. In the response you will get a nextCursor value, which can be used to get the next set of results. For the first call, this should be ignored."
},
"limit": {
"type": "integer",
"description": "Maximum number of results to return. Default is 10."
}
},
"required": [
"entityType"
]
}
},
{
"name": "create_test_case",
"description": "This tool can be used to create a test case for a table or a table's column. It needs test definitions. These test definition can be found using the get_test_definitions tool. If this fails do not do other operations.",
"parameters": {
"description": "This tool can be used to create a test case for a table or a table's column. It requires the test definition and the entity (table or column) to be tested.",
"type": "object",
"properties": {
"name": {
"type": "string",
"description": "Name of the test case. Either the user tells the name or auto generate a name based on the test definition and entity fqn."
},
"entityFqn": {
"type": "string",
"description": "Name of the asset or entity for which the test case is being added. It is asset fqn that user originally ask to add test for , and not column fqn in case of column tests."
},
"entityType": {
"type": "string",
"description": "Type of the entity for which the test case is being added. Default is 'table'. It can only be one of these values 'container', 'workflowDefinition', 'testSuite', 'pipelineService', 'messagingService', 'type', 'dataContract', 'metadataService', 'appMarketPlaceDefinition', 'tag', 'dashboard', 'app', 'dataProduct', 'persona', 'workflow', 'kpi', 'apiService', 'query', 'searchIndex', 'classification', 'dashboardDataModel', 'glossary', 'apiEndpoint', 'storageService', 'domain', 'topic', 'databaseSchema', 'role', 'dataInsightChart', 'bot', 'mlmodelService', 'document', 'ingestionPipeline', 'database', 'searchService', 'testConnectionDefinition', 'webAnalyticEvent', 'table', 'policy', 'storedProcedure', 'databaseService', 'eventsubscription', 'dashboardService', 'apiCollection', 'team', 'mlmodel', 'glossaryTerm', 'pipeline', 'dataInsightCustomChart', 'metric', 'report', 'chart', 'user', 'testCase', 'testDefinition'."
},
"testDefinitionName": {
"type": "string",
"description": "Fully qualified name of the test definition to be used for the test case."
},
"description": {
"type": "string",
"description": "Description of the test case. This can be auto generated based on entity details, the test definition and type of test."
},
"parameterValues": {
"type": "array",
"description": "Parameter values for the test case. This is an array of objects, where each object contains the parameter name (this is 'name' from parameter definition in test definition) and value (can be asked from user or auto supplied). These parameter can be found from the testDefinitions from `parameterDefinition` which has name, description, dataType and whether it is required mandatorily.",
"items": {
"type": "object",
"description": "Parameter values that can be passed for a Test Case.",
"properties": {
"name": {
"description": "Ma,e of the parameter. Must match the parameter names in testCaseParameterDefinition",
"type": "string"
},
"value": {
"description": "value to be passed for the Parameters. These are input from Users. We capture this in string and convert during the runtime.",
"type": "string"
}
},
"required": [
"name",
"value"
]
}
}
},
"required": [
"name",
"testDefinition",
"parameterValues"
]
}
}
]
}