From 4ac5912d4ccc883a617832a8f255c450c610fe3c Mon Sep 17 00:00:00 2001
From: Imri Paran <imri.paran@gmail.com>
Date: Fri, 26 Apr 2024 11:49:08 +0200
Subject: [PATCH] MINOR: added TestCase inspection query to backend and sample
 data (#16003)

* added TestCase inspection query to backend and sample data

* format

* format
---
 .../sample_data/tests/testCaseResults.json    |  3 +-
 .../ingestion/ometa/mixins/tests_mixin.py     | 15 +++++++++
 .../ingestion/source/database/sample_data.py  |  5 +++
 .../service/jdbi3/TestCaseRepository.java     | 12 +++++++
 .../resources/dqtests/TestCaseResource.java   | 28 +++++++++++++++++
 .../dqtests/TestCaseResourceTest.java         | 31 ++++++++++++++++++-
 .../resources/json/schema/tests/testCase.json |  4 +++
 7 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/ingestion/examples/sample_data/tests/testCaseResults.json b/ingestion/examples/sample_data/tests/testCaseResults.json
index c55d10c95ce..ea0794b8e95 100644
--- a/ingestion/examples/sample_data/tests/testCaseResults.json
+++ b/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;"
         }
     ]
 }
\ No newline at end of file
diff --git a/ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py b/ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py
index 3bddfe7d1d1..a7fa1ce3609 100644
--- a/ingestion/src/metadata/ingestion/ometa/mixins/tests_mixin.py
+++ b/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)
diff --git a/ingestion/src/metadata/ingestion/source/database/sample_data.py b/ingestion/src/metadata/ingestion/source/database/sample_data.py
index 73ecfafdee3..783e66c9f23 100644
--- a/ingestion/src/metadata/ingestion/source/database/sample_data.py
+++ b/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"""
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestCaseRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestCaseRepository.java
index 262c74b2566..02cbfb448ad 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestCaseRepository.java
+++ b/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(),
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestCaseResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestCaseResource.java
index 799c357b17b..5468b8dd249 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestCaseResource.java
+++ b/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(
diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/dqtests/TestCaseResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/dqtests/TestCaseResourceTest.java
index ffd80d16379..37f04d741c1 100644
--- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/dqtests/TestCaseResourceTest.java
+++ b/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");
diff --git a/openmetadata-spec/src/main/resources/json/schema/tests/testCase.json b/openmetadata-spec/src/main/resources/json/schema/tests/testCase.json
index 7c08dc10b70..f0b1b27da18 100644
--- a/openmetadata-spec/src/main/resources/json/schema/tests/testCase.json
+++ b/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"