From e4c3aba7537997d6b40a680c0a211170f59c69b9 Mon Sep 17 00:00:00 2001 From: Avinder Bahra Date: Thu, 27 Jun 2024 14:03:10 +0100 Subject: [PATCH] fix restrictions on partition and sort key expressions for complex cases (#444) --- .../scala/zio/dynamodb/PartitionKey.scala | 22 ++- .../src/main/scala/zio/dynamodb/SortKey.scala | 187 +++++++++++------- .../proofs/CanSortKeyBeginsWith.scala | 8 +- .../examples/KeyConditionExprExample.scala | 60 +++++- 4 files changed, 196 insertions(+), 81 deletions(-) diff --git a/dynamodb/src/main/scala/zio/dynamodb/PartitionKey.scala b/dynamodb/src/main/scala/zio/dynamodb/PartitionKey.scala index 6ca95e52..694ee2ca 100644 --- 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 index 7c25b3b5..fe0348a6 100644 --- 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 index c421410b..370ff03d 100644 --- 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 index 300f865b..3dec4260 100644 --- 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 + }