fix restrictions on partition and sort key expressions for complex ca…

…ses (#444)
zio · Jun 27, 2024 · e4c3aba · e4c3aba
1 parent ef3b66a
commit e4c3aba
Show file tree

Hide file tree

Showing 4 changed files with 196 additions and 81 deletions.
diff --git a/dynamodb/src/main/scala/zio/dynamodb/PartitionKey.scala b/dynamodb/src/main/scala/zio/dynamodb/PartitionKey.scala
@@ -1,13 +1,21 @@
 package zio.dynamodb
 
 import zio.dynamodb.KeyConditionExpr.PartitionKeyEquals
-import zio.dynamodb.proofs.RefersTo
+import zio.dynamodb.ProjectionExpression.Unknown
 
-private[dynamodb] final case class PartitionKey[-From, +To](keyName: String) { self =>
-  def ===[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    value: To2
-  )(implicit ev: RefersTo[To1, To2]): PartitionKeyEquals[From] = {
-    val _ = ev
-    PartitionKeyEquals(self, implicitly[ToAttributeValue[To2]].toAttributeValue(value))
+private[dynamodb] final case class PartitionKey[-From, +To](keyName: String)
+private[dynamodb] object PartitionKey {
+  implicit class PartitionKeyUnknownToOps[-From](val pk: PartitionKey[From, Unknown])         {
+    def ===[To: ToAttributeValue](
+      value: To
+    ): PartitionKeyEquals[From] =
+      PartitionKeyEquals(pk, implicitly[ToAttributeValue[To]].toAttributeValue(value))
   }
+  implicit class PartitionKeyOps[-From, To: ToAttributeValue](val pk: PartitionKey[From, To]) {
+    def ===(
+      value: To
+    ): PartitionKeyEquals[From] =
+      PartitionKeyEquals(pk, implicitly[ToAttributeValue[To]].toAttributeValue(value))
+  }
+
 }
diff --git a/dynamodb/src/main/scala/zio/dynamodb/SortKey.scala b/dynamodb/src/main/scala/zio/dynamodb/SortKey.scala
@@ -1,84 +1,129 @@
 package zio.dynamodb
 
-import zio.dynamodb.proofs.RefersTo
 import zio.dynamodb.proofs.CanSortKeyBeginsWith
 
 import zio.dynamodb.KeyConditionExpr.SortKeyEquals
 import zio.dynamodb.KeyConditionExpr.ExtendedSortKeyExpr
+import zio.dynamodb.ProjectionExpression.Unknown
 
-private[dynamodb] final case class SortKey[-From, +To](keyName: String) { self =>
+private[dynamodb] final case class SortKey[-From, +To](keyName: String)
+
+private[dynamodb] object SortKey {
   // all comparison ops apply to: Strings, Numbers, Binary values
-  def ===[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    value: To2
-  )(implicit ev: RefersTo[To1, To2]): SortKeyEquals[From] = {
-    val _ = ev
-    SortKeyEquals[From](
-      self.asInstanceOf[SortKey[From, To2]],
-      implicitly[ToAttributeValue[To2]].toAttributeValue(value)
-    )
-  }
-  def >[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    value: To2
-  )(implicit ev: RefersTo[To1, To2]): ExtendedSortKeyExpr[From, To2] = {
-    val _ = ev
-    ExtendedSortKeyExpr.GreaterThan(
-      self.asInstanceOf[SortKey[From, To2]],
-      implicitly(ToAttributeValue[To2]).toAttributeValue(value)
-    )
-  }
-  def <[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    value: To2
-  )(implicit ev: RefersTo[To1, To2]): ExtendedSortKeyExpr[From, To2] = {
-    val _ = ev
-    ExtendedSortKeyExpr.LessThan(
-      self.asInstanceOf[SortKey[From, To2]],
-      implicitly[ToAttributeValue[To2]].toAttributeValue(value)
-    )
-  }
-  def <>[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    value: To2
-  )(implicit ev: RefersTo[To1, To2]): ExtendedSortKeyExpr[From, To2] = {
-    val _ = ev
-    ExtendedSortKeyExpr.NotEqual(
-      self.asInstanceOf[SortKey[From, To2]],
-      implicitly(ToAttributeValue[To2]).toAttributeValue(value)
-    )
-  }
-  def <=[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    value: To2
-  )(implicit ev: RefersTo[To1, To2]): ExtendedSortKeyExpr[From, To2] = {
-    val _ = ev
-    ExtendedSortKeyExpr.LessThanOrEqual(
-      self.asInstanceOf[SortKey[From, To2]],
-      implicitly[ToAttributeValue[To2]].toAttributeValue(value)
-    )
-  }
-  def >=[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    value: To2
-  )(implicit ev: RefersTo[To1, To2]): ExtendedSortKeyExpr[From, To2] = {
-    val _ = ev
-    ExtendedSortKeyExpr.GreaterThanOrEqual(
-      self.asInstanceOf[SortKey[From, To2]],
-      implicitly[ToAttributeValue[To2]].toAttributeValue(value)
-    )
+
+  implicit class SortKeyUnknownToOps[-From](val sk: SortKey[From, Unknown]) {
+    def ===[To: ToAttributeValue](
+      value: To
+    ): SortKeyEquals[From]                                                             =
+      SortKeyEquals(sk, implicitly[ToAttributeValue[To]].toAttributeValue(value))
+    def >[To: ToAttributeValue](
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                                                   =
+      ExtendedSortKeyExpr.GreaterThan(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly(ToAttributeValue[To]).toAttributeValue(value)
+      )
+    def <[To: ToAttributeValue](
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                                                   =
+      ExtendedSortKeyExpr.LessThan(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(value)
+      )
+    def <>[To: ToAttributeValue](
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                                                   =
+      ExtendedSortKeyExpr.NotEqual(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly(ToAttributeValue[To]).toAttributeValue(value)
+      )
+    def <=[To: ToAttributeValue](
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                                                   =
+      ExtendedSortKeyExpr.LessThanOrEqual(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(value)
+      )
+    def >=[To: ToAttributeValue](
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                                                   =
+      ExtendedSortKeyExpr.GreaterThanOrEqual(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(value)
+      )
+    def between[To: ToAttributeValue](min: To, max: To): ExtendedSortKeyExpr[From, To] =
+      ExtendedSortKeyExpr.Between[From, To](
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(min),
+        implicitly[ToAttributeValue[To]].toAttributeValue(max)
+      )
+    def beginsWith[To: ToAttributeValue](
+      prefix: To
+    )(implicit ev: CanSortKeyBeginsWith[_, To]): ExtendedSortKeyExpr[From, To] = {
+      val _ = ev
+      ExtendedSortKeyExpr.BeginsWith[From, To](
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(prefix)
+      )
+    }
   }
-  // applies to all PK types
-  def between[To: ToAttributeValue, IsPrimaryKey](min: To, max: To): ExtendedSortKeyExpr[From, To] =
-    ExtendedSortKeyExpr.Between[From, To](
-      self.asInstanceOf[SortKey[From, To]],
-      implicitly[ToAttributeValue[To]].toAttributeValue(min),
-      implicitly[ToAttributeValue[To]].toAttributeValue(max)
-    )
 
-  // beginsWith applies to: Strings, Binary values
-  def beginsWith[To1 >: To, To2: ToAttributeValue, IsPrimaryKey](
-    prefix: To2
-  )(implicit ev: CanSortKeyBeginsWith[To1, To2]): ExtendedSortKeyExpr[From, To2] = {
-    val _ = ev
-    ExtendedSortKeyExpr.BeginsWith[From, To2](
-      self.asInstanceOf[SortKey[From, To2]],
-      implicitly[ToAttributeValue[To2]].toAttributeValue(prefix)
-    )
+  implicit class SortKeyOps[-From, To: ToAttributeValue](val sk: SortKey[From, To]) {
+    def ===(
+      value: To
+    ): SortKeyEquals[From]                                       =
+      SortKeyEquals(sk, implicitly[ToAttributeValue[To]].toAttributeValue(value))
+    def >(
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                             =
+      ExtendedSortKeyExpr.GreaterThan(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly(ToAttributeValue[To]).toAttributeValue(value)
+      )
+    def <(
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                             =
+      ExtendedSortKeyExpr.LessThan(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(value)
+      )
+    def <>(
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                             =
+      ExtendedSortKeyExpr.NotEqual(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly(ToAttributeValue[To]).toAttributeValue(value)
+      )
+    def <=(
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                             =
+      ExtendedSortKeyExpr.LessThanOrEqual(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(value)
+      )
+    def >=(
+      value: To
+    ): ExtendedSortKeyExpr[From, To]                             =
+      ExtendedSortKeyExpr.GreaterThanOrEqual(
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(value)
+      )
+    def between(min: To, max: To): ExtendedSortKeyExpr[From, To] =
+      ExtendedSortKeyExpr.Between[From, To](
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(min),
+        implicitly[ToAttributeValue[To]].toAttributeValue(max)
+      )
+    def beginsWith(
+      prefix: To
+    )(implicit ev: CanSortKeyBeginsWith[To, To]): ExtendedSortKeyExpr[From, To] = {
+      val _ = ev
+      ExtendedSortKeyExpr.BeginsWith[From, To](
+        sk.asInstanceOf[SortKey[From, To]],
+        implicitly[ToAttributeValue[To]].toAttributeValue(prefix)
+      )
+    }
+
   }
 
 }
diff --git a/dynamodb/src/main/scala/zio/dynamodb/proofs/CanSortKeyBeginsWith.scala b/dynamodb/src/main/scala/zio/dynamodb/proofs/CanSortKeyBeginsWith.scala
@@ -13,9 +13,13 @@ trait CanSortKeyBeginsWith0 extends CanSortKeyBeginsWith1 {
 }
 trait CanSortKeyBeginsWith1 {
   // begins_with with only applies to keys of type string or bytes
-  implicit def bytes[A <: Iterable[Byte]]: CanSortKeyBeginsWith[A, A] =
+  implicit def bytes[A <: Iterable[Byte]]: CanSortKeyBeginsWith[A, A]               =
     new CanSortKeyBeginsWith[A, A] {}
-  implicit def string: CanSortKeyBeginsWith[String, String]           = new CanSortKeyBeginsWith[String, String] {}
+  implicit def optionBytes[A <: Option[Iterable[Byte]]]: CanSortKeyBeginsWith[A, A] =
+    new CanSortKeyBeginsWith[A, A] {}
+  implicit def string: CanSortKeyBeginsWith[String, String]                         = new CanSortKeyBeginsWith[String, String] {}
+  implicit def optionString: CanSortKeyBeginsWith[Option[String], Option[String]]   =
+    new CanSortKeyBeginsWith[Option[String], Option[String]] {}
 }
 object CanSortKeyBeginsWith extends CanSortKeyBeginsWith0 {
   implicit def unknownLeft[X]: CanSortKeyBeginsWith[ProjectionExpression.Unknown, X] =

diff --git a/examples/src/main/scala/zio/dynamodb/examples/KeyConditionExprExample.scala b/examples/src/main/scala/zio/dynamodb/examples/KeyConditionExprExample.scala
@@ -4,6 +4,7 @@ import zio.dynamodb.ProjectionExpression
 
 import zio.schema.Schema
 import zio.schema.DeriveSchema
+import zio.schema.annotation.discriminatorName
 import zio.dynamodb.DynamoDBQuery
 
 object KeyConditionExprExample extends App {
@@ -51,10 +52,67 @@ object KeyConditionExprExample extends App {
   val pkAndSkExtended5 =
     Student.email.partitionKey === "x" && Student.binary.sortKey.beginsWith(List(1.toByte))
   val pkAndSkExtended6 =
-    Student.email.partitionKey === "x" && Student.binary2.sortKey.beginsWith(List(1.toByte))
+    Student.email.partitionKey === "x" && Student.binary2.sortKey.beginsWith(Vector(1.toByte))
 
   val (aliasMap, s) = pkAndSkExtended1.render.execute
   println(s"aliasMap=$aliasMap, s=$s")
 
   val get = DynamoDBQuery.queryAllItem("table").whereKey($("foo.bar").partitionKey === 1 && $("foo.baz").sortKey > 1)
+
+  // more complex example
+  @discriminatorName("invoiceType")
+  sealed trait InvoiceWithDiscriminatorName {
+    def id: String
+  }
+  object InvoiceWithDiscriminatorName       {
+    final case class Unpaid(id: String) extends InvoiceWithDiscriminatorName
+    object Unpaid {
+      implicit val schema: Schema.CaseClass1[String, Unpaid] = DeriveSchema.gen[Unpaid]
+      val id                                                 = ProjectionExpression.accessors[Unpaid]
+    }
+    final case class Paid(
+      id: String,
+      amount: Option[Int],
+      accountId: Option[String] = None,
+      email: Option[String] = None
+    ) extends InvoiceWithDiscriminatorName
+    object Paid   {
+      implicit val schema: Schema.CaseClass4[String, Option[Int], Option[String], Option[String], Paid] =
+        DeriveSchema.gen[Paid]
+      val (id, amount, accountId, email)                                                                = ProjectionExpression.accessors[Paid]
+    }
+    implicit val schema: Schema.Enum2[Unpaid, Paid, InvoiceWithDiscriminatorName] =
+      DeriveSchema.gen[InvoiceWithDiscriminatorName]
+    val (unpaid, paid) = ProjectionExpression.accessors[InvoiceWithDiscriminatorName]
+  }
+
+  val peId: ProjectionExpression[InvoiceWithDiscriminatorName, String]                =
+    InvoiceWithDiscriminatorName.paid >>> InvoiceWithDiscriminatorName.Paid.id
+  val peAccountId: ProjectionExpression[InvoiceWithDiscriminatorName, Option[String]] =
+    InvoiceWithDiscriminatorName.paid >>> InvoiceWithDiscriminatorName.Paid.accountId
+  val peEmailId: ProjectionExpression[InvoiceWithDiscriminatorName, Option[String]]   =
+    InvoiceWithDiscriminatorName.paid >>> InvoiceWithDiscriminatorName.Paid.email
+  val peAmount: ProjectionExpression[InvoiceWithDiscriminatorName, Option[Int]]       =
+    InvoiceWithDiscriminatorName.paid >>> InvoiceWithDiscriminatorName.Paid.amount
+  val pk1: PartitionKeyEquals[InvoiceWithDiscriminatorName]                           = peId.partitionKey === "1"
+  val pk2: PartitionKeyEquals[InvoiceWithDiscriminatorName]                           = peAccountId.partitionKey === Some("1")
+  val pk3: CompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]                      =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey === Some("1")
+  val pk4: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]              =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey > Some("1")
+  val pk5: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]              =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey < Some("1")
+  val pk6: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]              =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey <> Some("1")
+  val pk7: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]              =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey <= Some("1")
+  val pk8: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]              =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey >= Some("1")
+  val pk9: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]              =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey.between(Some("1"), Some("3"))
+  val pk10: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]             =
+    peAccountId.partitionKey === Some("1") && peEmailId.sortKey.beginsWith(Some("1"))
+  // val pk11: ExtendedCompositePrimaryKeyExpr[InvoiceWithDiscriminatorName]             =
+  //   peAccountId.partitionKey === Some("1") && peAmount.sortKey.beginsWith(Some(1)) // does not compile as expected
+
 }