Feature/348 create checkpoint v2 (#407)

Author: salamonpavel

.github/workflows/build.yml

@@ -52,7 +52,7 @@ jobs:
       postgres:
         image: postgres:15
         env:
-          POSTGRES_PASSWORD: postgres
+          POSTGRES_PASSWORD: changeme
           POSTGRES_DB: atum_db
         options: >-
           --health-cmd pg_isready

agent/src/main/scala/za/co/absa/atum/agent/dispatcher/ConsoleDispatcher.scala

@@ -18,15 +18,7 @@ package za.co.absa.atum.agent.dispatcher
 
 import com.typesafe.config.Config
 import org.apache.spark.internal.Logging
-import za.co.absa.atum.model.dto.{
-  AdditionalDataDTO,
-  AdditionalDataItemDTO,
-  AdditionalDataPatchDTO,
-  AtumContextDTO,
-  CheckpointDTO,
-  PartitioningDTO,
-  PartitioningSubmitDTO
-}
+import za.co.absa.atum.model.dto._
 
 /**
  *  dispatcher useful for development, testing and debugging

agent/src/main/scala/za/co/absa/atum/agent/dispatcher/Dispatcher.scala

@@ -17,14 +17,7 @@
 package za.co.absa.atum.agent.dispatcher
 
 import com.typesafe.config.Config
-import za.co.absa.atum.model.dto.{
-  AdditionalDataDTO,
-  AdditionalDataPatchDTO,
-  AtumContextDTO,
-  CheckpointDTO,
-  PartitioningDTO,
-  PartitioningSubmitDTO
-}
+import za.co.absa.atum.model.dto._
 
 /**
  *  This class provides a contract for different dispatchers. It has a constructor foe eventual creation via reflection.

agent/src/main/scala/za/co/absa/atum/agent/dispatcher/HttpDispatcher.scala

@@ -22,6 +22,7 @@ import sttp.capabilities
 import sttp.client3._
 import sttp.model.Uri
 import sttp.client3.okhttp.OkHttpSyncBackend
+import sttp.model.StatusCode
 import za.co.absa.atum.agent.exception.AtumAgentException.HttpException
 import za.co.absa.atum.model.ApiPaths
 import za.co.absa.atum.model.dto._
@@ -33,7 +34,6 @@ class HttpDispatcher(config: Config) extends Dispatcher(config) with Logging {
 
   private val serverUrl: String = config.getString(UrlKey)
 
-  private val apiV1 = s"/${ApiPaths.Api}/${ApiPaths.V1}"
   private val apiV2 = s"/${ApiPaths.Api}/${ApiPaths.V2}"
 
   private val getPartitioningIdEndpoint = Uri.unsafeParse(s"$serverUrl$apiV2/${ApiPaths.V2Paths.Partitionings}")
@@ -64,8 +64,8 @@ class HttpDispatcher(config: Config) extends Dispatcher(config) with Logging {
 
     val response = backend.send(request)
 
-    response.code.code match {
-      case 404 => None
+    response.code match {
+      case StatusCode.NotFound => None
       case _ => Some(handleResponseBody(response).as[SingleSuccessResponse[PartitioningWithIdDTO]].data)
     }
   }
@@ -104,27 +104,38 @@ class HttpDispatcher(config: Config) extends Dispatcher(config) with Logging {
       )
       val req = commonAtumRequest.get(endpoint)
       val resp = backend.send(req)
-      handleResponseBody(resp)
-        .as[SingleSuccessResponse[AdditionalDataDTO.Data]]
-        .data
-        .map(item => item._1 -> item._2.map(_.value))
+      handleResponseBody(resp).as[MultiSuccessResponse[AdditionalDataItemV2DTO]]
     }
 
     AtumContextDTO(
       partitioning = newPartitioningWithIdDTO.partitioning,
       measures = measures,
-      additionalData = additionalData
+      additionalData = additionalData.data.map(item => item.key -> item.value).toMap
     )
   }
 
   override protected[agent] def saveCheckpoint(checkpoint: CheckpointDTO): Unit = {
-    val endpoint = Uri.unsafeParse(s"$serverUrl$apiV1/${ApiPaths.V1Paths.CreateCheckpoint}")
+    val partitioningId = getPartitioningId(checkpoint.partitioning)
+
+    val checkpointV2DTO = CheckpointV2DTO(
+      id = checkpoint.id,
+      name = checkpoint.name,
+      author = checkpoint.author,
+      measuredByAtumAgent = checkpoint.measuredByAtumAgent,
+      processStartTime = checkpoint.processStartTime,
+      processEndTime = checkpoint.processEndTime,
+      measurements = checkpoint.measurements,
+      properties = checkpoint.properties
+    )
+
+    val endpoint = Uri.unsafeParse(
+      s"$serverUrl$apiV2/${ApiPaths.V2Paths.Partitionings}/$partitioningId/${ApiPaths.V2Paths.Checkpoints}"
+    )
     val request = commonAtumRequest
       .post(endpoint)
-      .body(checkpoint.asJsonString)
+      .body(checkpointV2DTO.asJsonString)
 
     val response = backend.send(request)
-
     handleResponseBody(response)
   }
 
@@ -145,7 +156,14 @@ class HttpDispatcher(config: Config) extends Dispatcher(config) with Logging {
 
     val response = backend.send(request)
 
-    handleResponseBody(response).as[SingleSuccessResponse[AdditionalDataDTO]].data
+    val data: AdditionalDataDTO.Data = handleResponseBody(response).as[MultiSuccessResponse[AdditionalDataItemV2DTO]]
+      .data
+      .map( item => item.value match {
+        case Some(_) => item.key -> Some(AdditionalDataItemDTO(item.value.get, item.author))
+        case None => item.key -> None
+      }).toMap
+
+    AdditionalDataDTO(data)
   }
 
   private def handleResponseBody(response: Response[Either[String, String]]): String = {

agent/src/test/resources/reference.conf

@@ -15,7 +15,7 @@
 atum.dispatcher.type="console"
 
 # The REST API URI of the atum server
-# atum.dispatcher.http.url="http://localhost:8080"
+#atum.dispatcher.http.url="http://localhost:8080"
 
 # Maximum number of dispatch captures to keep in memory
 # atum.dispatcher.capture.capture-limit=1000 # 0 means no limit

agent/src/test/scala/za/co/absa/atum/agent/AgentServerCompatibilityTests.scala

@@ -61,28 +61,28 @@ class AgentServerCompatibilityTests extends DBTestSuite {
     val rdd = spark.sparkContext.parallelize(testDataForRDD)
     val df = spark.createDataFrame(rdd, testDataSchema)
 
-    // Atum Agent preparation - Agent configured to work with HTTP Dispatcher and service on localhost
-    val agent: AtumAgent = new AtumAgent {
-      override val dispatcher: HttpDispatcher = new HttpDispatcher(
-        ConfigFactory
-          .empty()
-          .withValue("atum.dispatcher.type", ConfigValueFactory.fromAnyRef("http"))
-          .withValue("atum.dispatcher.http.url", ConfigValueFactory.fromAnyRef("http://localhost:8080"))
-      )
-    }
-
-    // Atum Context stuff preparation - Partitioning, Measures, Additional Data, Checkpoint
-    val domainAtumPartitioning = ListMap(
-      "partition1" -> "valueFromTest1",
-      "partition2" -> "valueFromTest2"
-    )
-    val domainAtumContext = agent.getOrCreateAtumContext(domainAtumPartitioning)
-
-    domainAtumContext.addMeasure(RecordCount())
-    domainAtumContext.addAdditionalData("author", "Laco")
-    domainAtumContext.addAdditionalData(Map("author" -> "LacoNew", "version" -> "1.0"))
-
-    domainAtumContext.createCheckpoint("checkPointNameCount", df)
+     // Atum Agent preparation - Agent configured to work with HTTP Dispatcher and service on localhost
+     val agent: AtumAgent = new AtumAgent {
+       override val dispatcher: HttpDispatcher = new HttpDispatcher(
+         ConfigFactory
+           .empty()
+           .withValue("atum.dispatcher.type", ConfigValueFactory.fromAnyRef("http"))
+           .withValue("atum.dispatcher.http.url", ConfigValueFactory.fromAnyRef("http://localhost:8080"))
+       )
+     }
+
+     // Atum Context stuff preparation - Partitioning, Measures, Additional Data, Checkpoint
+     val domainAtumPartitioning = ListMap(
+       "partition1" -> "valueFromTest1",
+       "partition2" -> "valueFromTest2"
+     )
+     val domainAtumContext = agent.getOrCreateAtumContext(domainAtumPartitioning)
+
+     domainAtumContext.addMeasure(RecordCount())
+     domainAtumContext.addAdditionalData("author", "Laco")
+     domainAtumContext.addAdditionalData(Map("author" -> "LacoNew", "version" -> "1.0"))
+
+     domainAtumContext.createCheckpoint("checkPointNameCount", df)
 
     // DB Check, data should be written in the DB
     table("runs.partitionings").all() { partitioningsResult =>
@@ -94,35 +94,45 @@ class AgentServerCompatibilityTests extends DBTestSuite {
     }
 
     table("runs.additional_data").all() { adResult =>
+      val expectedMap = Map(
+        "author" -> "LacoNew",
+        "version" -> "1.0"
+      )
+
       assert(adResult.hasNext)
       val row = adResult.next()
 
-      assert(row.getString("ad_name").contains("author"))
-      assert(row.getString("ad_value").contains("LacoNew"))
+      val adName1 = row.getString("ad_name").get
+      val adValue1 = row.getString("ad_value").get
+
+      assert(expectedMap(adName1) == adValue1)
 
       assert(adResult.hasNext)
       val row2 = adResult.next()
 
-      assert(row2.getString("ad_name").contains("version"))
-      assert(row2.getString("ad_value").contains("1.0"))
+      val adName2 = row2.getString("ad_name").get
+      val adValue2 = row2.getString("ad_value").get
+
+      assert(expectedMap(adName2) == adValue2)
+
       assert(!adResult.hasNext)
     }
 
-    table("runs.additional_data_history").all() { adHistResult =>
+  table("runs.additional_data_history").all() { adHistResult =>
       assert(adHistResult.hasNext)
       val row = adHistResult.next()
 
       assert(row.getString("ad_name").contains("author"))
       assert(row.getString("ad_value").contains("Laco"))
 
-      assert(!adHistResult.hasNext)
+      assert(!adHistResult.hasNext) // failing when executed multiple times
     }
 
     table("runs.measure_definitions").all() { measureDefResult =>
       assert(measureDefResult.hasNext)
       val row = measureDefResult.next()
 
-      assert(row.getString("measure_name").contains("*"))
+      assert(row.getString("measure_name").contains("count"))
       assert(row.getString("measured_columns").contains("{}"))
 
       assert(!measureDefResult.hasNext)
@@ -135,7 +145,5 @@ class AgentServerCompatibilityTests extends DBTestSuite {
       assert(row.getJsonB("measurement_value").contains(expectedMeasurement))
       assert(!measurementsResult.hasNext)
     }
-
-    // TODO Truncate data potentially, Balta might support this in the near future
   }
 }

agent/src/test/scala/za/co/absa/atum/agent/dispatcher/HttpDispatcherUnitTests.scala

@@ -2,17 +2,20 @@ package za.co.absa.atum.agent.dispatcher
 
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
-import org.scalatest.BeforeAndAfterEach
+import org.scalatest.{BeforeAndAfterEach, Ignore}
 import org.mockito.Mockito._
 import org.mockito.ArgumentMatchers._
 import sttp.client3._
 import sttp.model.StatusCode
 import com.typesafe.config.Config
 import sttp.capabilities
+import za.co.absa.atum.agent.exception.AtumAgentException.HttpException
 import za.co.absa.atum.model.dto._
 import za.co.absa.atum.model.envelopes.SuccessResponse.{MultiSuccessResponse, SingleSuccessResponse}
 import za.co.absa.atum.model.utils.JsonSyntaxExtensions.JsonSerializationSyntax
 
+import java.time.ZonedDateTime
+
 class HttpDispatcherUnitTests extends AnyFlatSpec with Matchers with BeforeAndAfterEach {
 
   var mockBackend: SttpBackend[Identity, sttp.capabilities.WebSockets] = _
@@ -23,9 +26,9 @@ class HttpDispatcherUnitTests extends AnyFlatSpec with Matchers with BeforeAndAf
   val testPartitioningSubmitDTO = PartitioningSubmitDTO(testPartitioningDTO, None, "author")
   val createdPartitioningWithId = PartitioningWithIdDTO(123L, testPartitioningDTO, "author")
   val measures: Seq[MeasureDTO] = Seq(MeasureDTO("m1", Seq("c1")), MeasureDTO("m2", Seq("c2")))
-  val additionalData: Map[String, Option[AdditionalDataItemDTO]] = Map(
-    "key1" -> Some(AdditionalDataItemDTO("val1", "author1")),
-    "key2" -> None
+  val additionalData: Seq[AdditionalDataItemV2DTO] = Seq(
+    AdditionalDataItemV2DTO("key1", Some("val1"), "author1"),
+    AdditionalDataItemV2DTO("key2", None, "author2")
   )
 
   def encodedPartitioning(partitioning: Seq[PartitionDTO]): String =
@@ -100,7 +103,7 @@ class HttpDispatcherUnitTests extends AnyFlatSpec with Matchers with BeforeAndAf
     val measuresResponse =
       Response(Right(MultiSuccessResponse(measures).asJsonString): Either[String, String], StatusCode.Ok)
     val additionalDataResponse =
-      Response(Right(SingleSuccessResponse(additionalData).asJsonString): Either[String, String], StatusCode.Ok)
+      Response(Right(MultiSuccessResponse[AdditionalDataItemV2DTO](additionalData).asJsonString): Either[String, String], StatusCode.Ok)
 
     stubGetPartitioning(parentPartitioning, getParentResponse)
     stubGetPartitioning(testPartitioningDTO, getPartitioningResponse)
@@ -158,7 +161,7 @@ class HttpDispatcherUnitTests extends AnyFlatSpec with Matchers with BeforeAndAf
     val measuresResponse =
       Response(Right(MultiSuccessResponse(Seq.empty[MeasureDTO]).asJsonString): Either[String, String], StatusCode.Ok)
     val additionalDataResponse = Response(
-      Right(SingleSuccessResponse(Map.empty[String, Option[AdditionalDataItemDTO]]).asJsonString): Either[
+      Right(MultiSuccessResponse(Seq.empty[AdditionalDataItemV2DTO]).asJsonString): Either[
         String,
         String
       ],
@@ -177,4 +180,83 @@ class HttpDispatcherUnitTests extends AnyFlatSpec with Matchers with BeforeAndAf
     result.measures shouldBe empty
     result.additionalData shouldBe empty
   }
+
+  "saveCheckpoint" should "successfully save a checkpoint when partitioning exists" in {
+    val checkpoint = CheckpointDTO(
+      id = java.util.UUID.randomUUID(),
+      name = "cp",
+      author = "author",
+      measuredByAtumAgent = true,
+      processStartTime = ZonedDateTime.now(),
+      processEndTime = Some(ZonedDateTime.now()),
+      measurements = Set.empty,
+      partitioning = testPartitioningDTO
+    )
+    val partitioningWithId = PartitioningWithIdDTO(123L, testPartitioningDTO, "author")
+    val getPartitioningResponse = Response(
+      Right(SingleSuccessResponse(partitioningWithId).asJsonString): Either[String, String],
+      StatusCode.Ok
+    )
+    val postCheckpointResponse = Response(Right("{}"): Either[String, String], StatusCode.Created)
+
+    stubGetPartitioning(testPartitioningDTO, getPartitioningResponse)
+    when(
+      mockBackend.send(
+        argThat[Request[Either[String, String], capabilities.WebSockets]](
+          req => req.method.method == "POST" && req.uri.path.mkString.contains("checkpoints")
+        )
+      )
+    ).thenReturn(postCheckpointResponse)
+
+    noException should be thrownBy dispatcher.saveCheckpoint(checkpoint)
+  }
+
+  it should "throw HttpException if partitioning is not found" in {
+    val checkpoint = CheckpointDTO(
+      id = java.util.UUID.randomUUID(),
+      name = "cp",
+      author = "author",
+      measuredByAtumAgent = true,
+      processStartTime = ZonedDateTime.now(),
+      processEndTime = Some(ZonedDateTime.now()),
+      measurements = Set.empty,
+      partitioning = testPartitioningDTO
+    )
+    val notFoundResponse = Response(Left("Not found"): Either[String, String], StatusCode.NotFound)
+
+    stubGetPartitioning(testPartitioningDTO, notFoundResponse)
+
+    an[HttpException] should be thrownBy dispatcher.saveCheckpoint(checkpoint)
+  }
+
+  it should "throw HttpException if server returns error on checkpoint save" in {
+    val checkpoint = CheckpointDTO(
+      id = java.util.UUID.randomUUID(),
+      name = "cp",
+      author = "author",
+      measuredByAtumAgent = true,
+      processStartTime = ZonedDateTime.now(),
+      processEndTime = Some(ZonedDateTime.now()),
+      measurements = Set.empty,
+      partitioning = testPartitioningDTO
+    )
+    val partitioningWithId = PartitioningWithIdDTO(123L, testPartitioningDTO, "author")
+    val getPartitioningResponse = Response(
+      Right(SingleSuccessResponse(partitioningWithId).asJsonString): Either[String, String],
+      StatusCode.Ok
+    )
+    val badRequestResponse = Response(Left("Bad request"): Either[String, String], StatusCode.BadRequest)
+
+    stubGetPartitioning(testPartitioningDTO, getPartitioningResponse)
+    when(
+      mockBackend.send(
+        argThat[Request[Either[String, String], capabilities.WebSockets]](
+          req => req.method.method == "POST" && req.uri.path.mkString.contains("checkpoints")
+        )
+      )
+    ).thenReturn(badRequestResponse)
+
+    an[HttpException] should be thrownBy dispatcher.saveCheckpoint(checkpoint)
+  }
+
 }