put with narrow (#498)

docs/index.md

@@ -32,7 +32,7 @@ To use the new Cats Effect 3 interop module, we need to also add the following l
 
 ```scala
 libraryDependencies ++= Seq(
-  "dev.zio" %% "zio-dynamodb-ce" "@VERSION@"
+  "dev.zio" %% "zio-dynamodb-ce" % "@VERSION@"
 )
 ```
 
@@ -43,7 +43,7 @@ AWS tools like the CLI and Console read/write a special JSON representation of d
 
 ```scala
 libraryDependencies ++= Seq(
-  "dev.zio" %% "zio-dynamodb-json" "@VERSION@"
+  "dev.zio" %% "zio-dynamodb-json" % "@VERSION@"
 )
 ```
 

dynamodb/src/it/scala/zio/dynamodb/TypeSafeApiMappingSpec.scala

@@ -14,7 +14,7 @@ import zio.test.TestAspect
 
 object TypeSafeApiMappingSpec extends DynamoDBLocalSpec {
 
-  // note the default sum type mapping does not work with top level sum types as it required an intermediate map
+  // note the default sum type mapping does not work with top level sum types as it requires an intermediate map
   // and partition keys must be scalar values
 
   @discriminatorName("invoiceType")

dynamodb/src/it/scala/zio/dynamodb/TypeSafeApiNarrowSpec.scala

@@ -51,12 +51,17 @@ object TypeSafeApiNarrowSpec extends DynamoDBLocalSpec {
     ) @@ TestAspect.nondeterministic
 
   val topLevelSumTypeNarrowSuite = suite("for top level Invoice sum type with @discriminatorName annotation")(
-    test("getWithNarrow succeeds in narrowing an Unpaid Invoice instance to Unpaid") {
+    test("put with narrow") {
       withSingleIdKeyTable { invoiceTable =>
         val keyCond: KeyConditionExpr.PartitionKeyEquals[dynamo.Invoice.Unpaid] =
           dynamo.Invoice.Unpaid.id.partitionKey === "1"
         for {
-          _    <- put[dynamo.Invoice](invoiceTable, dynamo.Invoice.Unpaid("1")).execute
+          _    <- DynamoDBQuery
+                    .putWithNarrow[dynamo.Invoice, dynamo.Invoice.Unpaid](invoiceTable, dynamo.Invoice.Unpaid("1"))
+                    .where(
+                      !dynamo.Invoice.Unpaid.id.exists
+                    ) // note expressions are of concrete type Unpaid eg ConditionExpression[dynamo.Invoice.Unpaid]
+                    .execute
           item <- getItem(invoiceTable, PrimaryKey("id" -> "1")).execute
 
           unpaid <- getWithNarrow[dynamo.Invoice, dynamo.Invoice.Unpaid](invoiceTable)(keyCond).execute.absolve
@@ -67,19 +72,31 @@ object TypeSafeApiNarrowSpec extends DynamoDBLocalSpec {
         }
       }
     },
+    test("getWithNarrow succeeds in narrowing an Unpaid Invoice instance to Unpaid") {
+      withSingleIdKeyTable { invoiceTable =>
+        val keyCond: KeyConditionExpr.PartitionKeyEquals[dynamo.Invoice.Unpaid] =
+          dynamo.Invoice.Unpaid.id.partitionKey === "1"
+        for {
+          _ <- put[dynamo.Invoice](invoiceTable, dynamo.Invoice.Unpaid("1")).execute
+
+          unpaid <- getWithNarrow[dynamo.Invoice, dynamo.Invoice.Unpaid](invoiceTable)(keyCond).execute.absolve
+        } yield {
+          val unpaid2: dynamo.Invoice.Unpaid = unpaid
+          assertTrue(unpaid2 == dynamo.Invoice.Unpaid("1"))
+        }
+      }
+    },
     test("getWithNarrow succeeds in narrowing an Paid Invoice instance to Paid") {
       withSingleIdKeyTable { invoiceTable =>
         val keyCond: KeyConditionExpr.PartitionKeyEquals[dynamo.Invoice.Paid] =
           dynamo.Invoice.Paid.id.partitionKey === "1"
         for {
-          _    <- put[dynamo.Invoice](invoiceTable, dynamo.Invoice.Paid("1", 42)).execute
-          item <- getItem(invoiceTable, PrimaryKey("id" -> "1")).execute
+          _ <- put[dynamo.Invoice](invoiceTable, dynamo.Invoice.Paid("1", 42)).execute
 
           paid <- getWithNarrow[dynamo.Invoice, dynamo.Invoice.Paid](invoiceTable)(keyCond).execute.absolve
         } yield {
           val paid2: dynamo.Invoice.Paid = paid
-          val ensureDiscriminatorPresent = item == Some(Item("id" -> "1", "invoiceType" -> "Paid", "amount" -> 42))
-          assertTrue(paid2 == dynamo.Invoice.Paid("1", 42) && ensureDiscriminatorPresent)
+          assertTrue(paid2 == dynamo.Invoice.Paid("1", 42))
         }
       }
     },
@@ -95,7 +112,7 @@ object TypeSafeApiNarrowSpec extends DynamoDBLocalSpec {
         )
       }
     },
-    test("getWithNarrow fails in narrowing an Paid Invoice instance to Unpaid") {
+    test("getWithNarrow fails in narrowing a Paid Invoice instance to Unpaid") {
       withSingleIdKeyTable { invoiceTable =>
         val keyCond: KeyConditionExpr.PartitionKeyEquals[dynamo.Invoice.Unpaid] =
           dynamo.Invoice.Unpaid.id.partitionKey === "1"

dynamodb/src/main/scala/zio/dynamodb/DynamoDBQuery.scala

@@ -481,10 +481,14 @@ object DynamoDBQuery {
     get(tableName, primaryKeyExpr.asAttrMap, ProjectionExpression.projectionsFromSchema[From])
 
   /**
-   * It is common practice to save top level sum types to DynamoDB and often we want to retrieve them back as the subtype.
+   * It is common practice to save top level sum types to DynamoDB and often we want to retrieve them back as the subtype
+   * with expressions in terms of the subtype as well.
    * `getWithNarrow` does a `get` with a safe narrow operation from type `From` to `To`.
    * If the narrow fails it returns a Decoding error with details of the cast failure in the message.
+   *
    * Requires implicit schemas in scope which ensure that `From` is an enum (sealed trait) and `To` is a record (case class) subtype.
+   *
+   * Note this is an experimental API and may be subject to change.
    */
   def getWithNarrow[From: Schema.Enum, To <: From: Schema.Record](tableName: String)(
     primaryKeyExpr: KeyConditionExpr.PrimaryKeyExpr[To]
@@ -503,8 +507,10 @@ object DynamoDBQuery {
     }
   }
 
-  // Safely narrows `a: From` to subtype type `To` and requires that there are implicit schemas in scope which
-  // ensure that `From` is an enum (sealed trait) and `To` is a record (case class) subtype.
+  /**
+   * Safely narrows `a: From` to subtype type `To` and requires that there are implicit schemas in scope which
+   * ensure that `From` is an enum (sealed trait) and `To` is a record (case class) subtype.
+   */
   private[dynamodb] def narrow[From: Schema.Enum, To <: From: Schema.Record](
     a: From
   ): Either[String, To] = {
@@ -551,6 +557,26 @@ object DynamoDBQuery {
   def put[A: Schema](tableName: String, a: A): DynamoDBQuery[A, Option[A]] =
     putItem(tableName, toItem(a)).map(_.flatMap(item => fromItem(item).toOption))
 
+  /**
+   * It is common to save the top level sum type to DynamoDB and often we want to save them back as the subtype
+   * with expressions in terms of the subtype as well.
+   * `putWithNarrow` does a `put` of type `To` which is widened to type `From` before the save to ensure that the discriminator is saved,
+   * and narrows the returned DynamoDBQuery to `To` .
+   *
+   * Requires implicit schemas in scope which ensure that `From` is an enum (sealed trait) and `To` is a record (case class) subtype.
+   *
+   * Note this is an experimental API and may be subject to change.
+   */
+  def putWithNarrow[From: Schema.Enum, To <: From: Schema.Record](
+    tableName: String,
+    a: To
+  ): DynamoDBQuery[To, Option[To]] = {
+    val fromEnumSchema = implicitly[Schema.Enum[From]]
+    val toSchema       = implicitly[Schema.Record[To]]
+    putItem(tableName, toItem(a.asInstanceOf[From])(fromEnumSchema))
+      .map(_.flatMap(item => fromItem(item)(toSchema).toOption))
+  }
+
   private[dynamodb] def toItem[A](a: A)(implicit schema: Schema[A]): Item =
     FromAttributeValue.attrMapFromAttributeValue
       .fromAttributeValue(AttributeValue.encode(a)(schema))