MINOR: added TestCase inspection query to backend and sample data (#16003)

* added TestCase inspection query to backend and sample data

* format

* format

ingestion/examples/sample_data/tests/testCaseResults.json

@@ -654,7 +654,8 @@
                                 "(830)112-9566x8681"
                             ]
                         ]
-                    }
+                    },
+            "inspectionQuery": "SELECT * FROM ecommerce_db.shopify.dim_address\nWHERE zip NOT BETWEEN 90001 AND 96162\nLIMIT 50;"
         }
     ]
 }

ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py

@@ -362,3 +362,18 @@ class OMetaTestsMixin:
                 )
 
         return None
+
+    def ingest_inspection_query(
+        self, test_case: TestCase, inspection_query: str
+    ) -> Optional[TestCase]:
+        """
+        PUT inspection query for a test case.
+
+        :param test_case: The test case that failed
+        :param inspection_query: SQL query to inspect the failed rows
+        """
+        resp = self.client.put(
+            f"{self.get_suffix(TestCase)}/{test_case.id.__root__}/inspectionQuery",
+            data=inspection_query,
+        )
+        return TestCase(**resp)

ingestion/src/metadata/ingestion/source/database/sample_data.py

@@ -1529,6 +1529,11 @@ class SampleDataSource(
                         columns=test_case_results["failedRowsSample"]["columns"],
                     ),
                 )
+            if test_case_results.get("inspectionQuery"):
+                self.metadata.ingest_inspection_query(
+                    case,
+                    test_case_results["inspectionQuery"],
+                )
 
     def ingest_data_insights(self) -> Iterable[Either[OMetaDataInsightSample]]:
         """Iterate over all the data insights and ingest them"""

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

@@ -745,6 +745,17 @@ public class TestCaseRepository extends EntityRepository<TestCase> {
     return testCase.withFailedRowsSample(tableData);
   }
 
+  @Transaction
+  public TestCase addInspectionQuery(UriInfo uri, UUID testCaseId, String sql) {
+    TestCase original = get(uri, testCaseId, getFields(PATCH_FIELDS));
+    TestCase updated =
+        JsonUtils.readValue(JsonUtils.pojoToJson(original), TestCase.class)
+            .withInspectionQuery(sql);
+    EntityUpdater entityUpdater = getUpdater(original, updated, Operation.PATCH);
+    entityUpdater.update();
+    return updated;
+  }
+
   @Transaction
   public RestUtil.DeleteResponse<TableData> deleteTestCaseFailedRowsSample(UUID id) {
     daoCollection.entityExtensionDAO().delete(id, FAILED_ROWS_SAMPLE_EXTENSION);
@@ -908,6 +919,7 @@ public class TestCaseRepository extends EntityRepository<TestCase> {
           TEST_CASE,
           updated.getId());
       recordChange("parameterValues", original.getParameterValues(), updated.getParameterValues());
+      recordChange("inspectionQuery", original.getInspectionQuery(), updated.getInspectionQuery());
       recordChange(
           "computePassedFailedRowCount",
           original.getComputePassedFailedRowCount(),

openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestCaseResource.java

@@ -981,6 +981,34 @@ public class TestCaseResource extends EntityResource<TestCase, TestCaseRepositor
     return addHref(uriInfo, repository.addFailedRowsSample(testCase, tableData));
   }
 
+  @PUT
+  @Path("/{id}/inspectionQuery")
+  @Operation(
+      operationId = "addInspectionQuery",
+      summary = "Add inspection query data",
+      description = "Add an inspection query for this test case.",
+      responses = {
+        @ApiResponse(
+            responseCode = "200",
+            description = "Successfully update the test case with an inspection query.",
+            content =
+                @Content(
+                    mediaType = "application/json",
+                    schema = @Schema(implementation = TestCase.class)))
+      })
+  public TestCase addInspectionQuery(
+      @Context UriInfo uriInfo,
+      @Context SecurityContext securityContext,
+      @Parameter(description = "Id of the test case", schema = @Schema(type = "UUID"))
+          @PathParam("id")
+          UUID id,
+      @Valid String query) {
+    OperationContext operationContext =
+        new OperationContext(entityType, MetadataOperation.EDIT_TESTS);
+    authorizer.authorize(securityContext, operationContext, getResourceContextById(id));
+    return addHref(uriInfo, repository.addInspectionQuery(uriInfo, id, query));
+  }
+
   @GET
   @Path("/{id}/failedRowsSample")
   @Operation(

openmetadata-service/src/test/java/org/openmetadata/service/resources/dqtests/TestCaseResourceTest.java

@@ -1853,7 +1853,7 @@ public class TestCaseResourceTest extends EntityResourceTest<TestCase, CreateTes
   public TestCase getTestCase(String fqn, Map<String, String> authHeaders)
       throws HttpResponseException {
     WebTarget target = getCollection().path("/name/" + fqn);
-    target = target.queryParam("fields", "incidentId");
+    target = target.queryParam("fields", "incidentId,inspectionQuery");
     return TestUtils.get(target, TestCase.class, authHeaders);
   }
 
@@ -2348,6 +2348,29 @@ public class TestCaseResourceTest extends EntityResourceTest<TestCase, CreateTes
             .count());
   }
 
+  @Test
+  void test_addInspectionQuery(TestInfo test) throws IOException {
+    CreateTestCase create =
+        createRequest(test)
+            .withEntityLink(TABLE_LINK)
+            .withTestSuite(TEST_SUITE1.getFullyQualifiedName())
+            .withTestDefinition(TEST_DEFINITION3.getFullyQualifiedName())
+            .withParameterValues(
+                List.of(
+                    new TestCaseParameterValue().withValue("100").withName("missingCountValue")));
+    TestCase testCase = createAndCheckEntity(create, ADMIN_AUTH_HEADERS);
+    String inspectionQuery = "SELECT * FROM test_table WHERE column1 = 'value1'";
+    putInspectionQuery(testCase, inspectionQuery, ADMIN_AUTH_HEADERS);
+    TestCase updated = getTestCase(testCase.getFullyQualifiedName(), ADMIN_AUTH_HEADERS);
+    assertEquals(updated.getInspectionQuery(), inspectionQuery);
+  }
+
+  private void putInspectionQuery(TestCase testCase, String sql, Map<String, String> authHeaders)
+      throws IOException {
+    TestCase putResponse = putInspectionQuery(testCase.getId(), sql, authHeaders);
+    assertEquals(sql, putResponse.getInspectionQuery());
+  }
+
   private void putFailedRowsSample(
       TestCase testCase,
       List<String> columns,
@@ -2375,6 +2398,12 @@ public class TestCaseResourceTest extends EntityResourceTest<TestCase, CreateTes
     return TestUtils.put(target, data, TestCase.class, OK, authHeaders);
   }
 
+  public TestCase putInspectionQuery(UUID testCaseId, String sql, Map<String, String> authHeaders)
+      throws HttpResponseException {
+    WebTarget target = getResource(testCaseId).path("/inspectionQuery");
+    return TestUtils.put(target, sql, TestCase.class, OK, authHeaders);
+  }
+
   public TableData getSampleData(UUID testCaseId, Map<String, String> authHeaders)
       throws HttpResponseException {
     WebTarget target = getResource(testCaseId).path("/failedRowsSample");

openmetadata-spec/src/main/resources/json/schema/tests/testCase.json

@@ -115,6 +115,10 @@
       "description": "Sample of failed rows for this test case.",
       "$ref": "../entity/data/table.json#/definitions/tableData"
     },
+    "inspectionQuery": {
+      "description": "SQL query to retrieve the failed rows for this test case.",
+      "$ref": "../type/basic.json#/definitions/sqlQuery"
+    },
     "domain" : {
       "description": "Domain the test case belongs to. When not set, the test case inherits the domain from the table it belongs to.",
       "$ref": "../type/entityReference.json"