narrowed get experiment

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

@@ -0,0 +1,65 @@
+package zio.dynamodb
+
+import zio.dynamodb.DynamoDBQuery.{ getItem, put }
+import zio.Scope
+import zio.test.Spec
+import zio.test.assertTrue
+import zio.test.TestEnvironment
+import zio.schema.Schema
+import zio.schema.DeriveSchema
+import zio.schema.annotation.discriminatorName
+import zio.test.TestAspect
+import zio.dynamodb.DynamoDBQuery.getWithNarrow
+
+object TypeSafeApiNarrowSpec extends DynamoDBLocalSpec {
+
+  object dynamo {
+    @discriminatorName("invoiceType")
+    sealed trait Invoice {
+      def id: String
+    }
+    object Invoice       {
+      final case class Unpaid(id: String) extends Invoice
+      object Unpaid {
+        implicit val schema: Schema.CaseClass1[String, Unpaid] = DeriveSchema.gen[Unpaid]
+        val id                                                 = ProjectionExpression.accessors[Unpaid]
+      }
+      final case class Paid(id: String, amount: Int) extends Invoice
+      object Paid   {
+        implicit val schema: Schema.CaseClass2[String, Int, Paid] = DeriveSchema.gen[Paid]
+        val (id, amount)                                          = ProjectionExpression.accessors[Paid]
+      }
+      implicit val schema: Schema.Enum2[Unpaid, Paid, Invoice] =
+        DeriveSchema.gen[Invoice]
+      val (unpaid, paid) = ProjectionExpression.accessors[Invoice]
+    }
+
+  }
+
+  override def spec: Spec[Environment with TestEnvironment with Scope, Any] =
+    suite("TypeSafeApiMappingSpec")(
+      topLevelSumTypeDiscriminatorNameSuite
+    ) @@ TestAspect.nondeterministic
+
+  val topLevelSumTypeDiscriminatorNameSuite = suite("with @discriminatorName annotation")(
+    test("getWithNarrow succeeds in narrowing an Invoice to Unpaid") {
+      withSingleIdKeyTable { invoiceTable =>
+        val key: ProjectionExpression[dynamo.Invoice, String]            = dynamo.Invoice.unpaid >>> dynamo.Invoice.Unpaid.id
+        val keyCond: KeyConditionExpr.PartitionKeyEquals[dynamo.Invoice] = key.partitionKey === "1"
+        for {
+          _      <- put[dynamo.Invoice](invoiceTable, dynamo.Invoice.Unpaid("1")).execute
+          unpaid <- getWithNarrow(dynamo.Invoice.unpaid)(invoiceTable)(keyCond).execute.absolve
+          item   <- getItem(invoiceTable, PrimaryKey("id" -> "1")).execute
+        } yield {
+          val unpaid2: dynamo.Invoice.Unpaid = unpaid
+          assertTrue(
+            unpaid2 == dynamo.Invoice.Unpaid("1") && item == Some(
+              Item("id" -> "1", "invoiceType" -> "Unpaid")
+            )
+          )
+        }
+      }
+    }
+  )
+
+}

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

@@ -480,6 +480,26 @@ object DynamoDBQuery {
   ): DynamoDBQuery[From, Either[ItemError, From]] =
     get(tableName, primaryKeyExpr.asAttrMap, ProjectionExpression.projectionsFromSchema[From])
 
+  /**
+   * When dealing with ADTs it is often necessary to narrow the type of the item returned from the database.
+   * `getWithNarrow` does a `get` with a narrow from type `From` to `To` using `narrowEvidence: ProjectionExpression[From, To]` as evidence that this is safe.
+   * If the case fails it returns a Decoding error
+   */
+  def getWithNarrow[From: Schema, To <: From](narrowEvidence: ProjectionExpression[From, To])(tableName: String)(
+    primaryKeyExpr: KeyConditionExpr.PrimaryKeyExpr[From]
+  ): DynamoDBQuery[From, Either[ItemError, To]] = {
+    val _ = narrowEvidence // provides derived schema guarantee that we can narrow
+    get(tableName)(primaryKeyExpr).map {
+      case Right(from) =>
+        scala.util.Try(from.asInstanceOf[To]) match { // this should be safe
+          case scala.util.Success(to) => Right(to)
+          case scala.util.Failure(_)  =>
+            Left(ItemError.DecodingError(s"failed to narrow from $from using evidence $narrowEvidence"))
+        }
+      case Left(error) => Left(error)
+    }
+  }
+
   private def get[A: Schema](
     tableName: String,
     key: PrimaryKey,