Series/2.x pk experiment 3 type params

dynamodb/src/main/scala/zio/dynamodb/DynamoDBExecutorImpl.scala

@@ -773,7 +773,7 @@ case object DynamoDBExecutorImpl {
   def awsQueryRequest(queryAll: QueryAll): QueryRequest = {
     val (aliasMap, (maybeFilterExpr, maybeKeyExpr, projections)) = (for {
       filter      <- AliasMapRender.collectAll(queryAll.filterExpression.map(_.render))
-      keyExpr     <- AliasMapRender.collectAll(queryAll.keyConditionExpression.map(_.render))
+      keyExpr     <- AliasMapRender.collectAll(queryAll.keyConditionExpr.map(_.render))
       projections <- AliasMapRender.forEach(queryAll.projections)
     } yield (filter, keyExpr, projections)).execute
 
@@ -797,7 +797,7 @@ case object DynamoDBExecutorImpl {
   private def awsQueryRequest(querySome: QuerySome): QueryRequest = {
     val (aliasMap, (maybeFilterExpr, maybeKeyExpr, projections)) = (for {
       filter      <- AliasMapRender.collectAll(querySome.filterExpression.map(_.render))
-      keyExpr     <- AliasMapRender.collectAll(querySome.keyConditionExpression.map(_.render))
+      keyExpr     <- AliasMapRender.collectAll(querySome.keyConditionExpr.map(_.render))
       projections <- AliasMapRender.forEach(querySome.projections)
     } yield (filter, keyExpr, projections)).execute
 

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

@@ -1,7 +1,7 @@
 package zio.dynamodb
 
 import zio.dynamodb.DynamoDBError.ValueNotFound
-import zio.dynamodb.proofs.{ CanFilter, CanWhere, CanWhereKey }
+import zio.dynamodb.proofs.{ CanFilter, CanWhere }
 import zio.dynamodb.DynamoDBQuery.BatchGetItem.TableGet
 import zio.dynamodb.DynamoDBQuery.BatchWriteItem.{ Delete, Put }
 import zio.dynamodb.DynamoDBQuery.{
@@ -28,6 +28,7 @@ import zio.prelude.ForEachOps
 import zio.schema.Schema
 import zio.stream.Stream
 import zio.{ Chunk, NonEmptyChunk, Schedule, ZIO, Zippable => _, _ }
+import zio.dynamodb.proofs.IsPrimaryKey
 
 sealed trait DynamoDBQuery[-In, +Out] { self =>
 
@@ -220,7 +221,7 @@ sealed trait DynamoDBQuery[-In, +Out] { self =>
   }
 
   /**
-   * Parallel executes a DynamoDB Scan in parallel.
+   * Executes a DynamoDB Scan in parallel.
    * There are no guarantees on order of returned items.
    *
    * @param n The number of parallel requests to make to DynamoDB
@@ -317,59 +318,29 @@ sealed trait DynamoDBQuery[-In, +Out] { self =>
   def selectCount: DynamoDBQuery[In, Out]                  = select(Select.Count)
 
   /**
-   * Adds a KeyConditionExpression to a DynamoDBQuery. Example:
+   * Adds a KeyConditionExpr to a DynamoDBQuery. Example:
    * {{{
-   * val newQuery = query.whereKey(partitionKey("email") === "avi@gmail.com" && sortKey("subject") === "maths")
+   * // high level type safe API where "email" and "subject" keys are defined using ProjectionExpression.accessors[Student]
+   * val newQuery = query.whereKey(email.partitionKey === "avi@gmail.com" && subject.sortKey === "maths")
+   *
+   * // low level API
+   * val newQuery = query.whereKey($("email").partitionKey === "avi@gmail.com" && $("subject").sortKey === "maths")
    * }}}
    */
-  def whereKey(keyConditionExpression: KeyConditionExpression): DynamoDBQuery[In, Out] =
+  def whereKey[From, Pk, Sk](keyConditionExpression: KeyConditionExpr[From, Pk, Sk]): DynamoDBQuery[In, Out] =
     self match {
       case Zip(left, right, zippable) =>
         Zip(left.whereKey(keyConditionExpression), right.whereKey(keyConditionExpression), zippable)
       case Map(query, mapper)         => Map(query.whereKey(keyConditionExpression), mapper)
       case Absolve(query)             => Absolve(query.whereKey(keyConditionExpression))
 
       case s: QuerySome =>
-        s.copy(keyConditionExpression = Some(keyConditionExpression)).asInstanceOf[DynamoDBQuery[In, Out]]
+        s.copy(keyConditionExpr = Some(keyConditionExpression)).asInstanceOf[DynamoDBQuery[In, Out]]
       case s: QueryAll  =>
-        s.copy(keyConditionExpression = Some(keyConditionExpression)).asInstanceOf[DynamoDBQuery[In, Out]]
+        s.copy(keyConditionExpr = Some(keyConditionExpression)).asInstanceOf[DynamoDBQuery[In, Out]]
       case _            => self
     }
 
-  /**
-   * Adds a KeyConditionExpression from a ConditionExpression to a DynamoDBQuery
-   * Must be in the form of `<Condition1> && <Condition2>` where format of `<Condition1>` is:
-   * {{{<ProjectionExpressionForPartitionKey> === <value>}}}
-   * and the format of `<Condition2>` is:
-   * {{{<ProjectionExpressionForSortKey> <op> <value>}}} where op can be one of `===`, `>`, `>=`, `<`, `<=`, `between`, `beginsWith`
-   *
-   * Example using type safe API:
-   * {{{
-   * // email and subject are partition and sort keys respectively
-   * val (email, subject, enrollmentDate, payment) = ProjectionExpression.accessors[Student]
-   * // ...
-   * val newQuery = query.whereKey(email === "avi@gmail.com" && subject === "maths")
-   * }}}
-   */
-  def whereKey[B](
-    conditionExpression: ConditionExpression[B]
-  )(implicit ev: CanWhereKey[B, Out]): DynamoDBQuery[In, Out] = {
-    val _                                              = ev
-    val keyConditionExpression: KeyConditionExpression =
-      KeyConditionExpression.fromConditionExpressionUnsafe(conditionExpression)
-    self match {
-      case Zip(left, right, zippable) =>
-        Zip(left.whereKey(keyConditionExpression), right.whereKey(keyConditionExpression), zippable)
-      case Map(query, mapper)         => Map(query.whereKey(keyConditionExpression), mapper)
-      case Absolve(query)             => Absolve(query.whereKey(keyConditionExpression))
-      case s: QuerySome               =>
-        s.copy(keyConditionExpression = Some(keyConditionExpression)).asInstanceOf[DynamoDBQuery[In, Out]]
-      case s: QueryAll                =>
-        s.copy(keyConditionExpression = Some(keyConditionExpression)).asInstanceOf[DynamoDBQuery[In, Out]]
-      case _                          => self
-    }
-  }
-
   def withRetryPolicy(retryPolicy: Schedule[Any, Throwable, Any]): DynamoDBQuery[In, Out] =
     self match {
       case Zip(left, right, zippable) =>
@@ -482,7 +453,17 @@ object DynamoDBQuery {
   ): DynamoDBQuery[Any, Option[Item]] =
     GetItem(TableName(tableName), key, projections.toList)
 
-  def get[A: Schema](
+  def get[From: Schema, Pk, Sk](
+    tableName: String,
+    projections: ProjectionExpression[_, _]*
+  )(
+    partitionKeyExpr: KeyConditionExpr.PrimaryKeyExpr[From, Pk, Sk]
+  )(implicit ev: IsPrimaryKey[Pk], ev2: IsPrimaryKey[Sk]): DynamoDBQuery[From, Either[DynamoDBError, From]] = {
+    val (_, _) = (ev, ev2)
+    get(tableName, partitionKeyExpr.asAttrMap, projections: _*)
+  }
+
+  private def get[A: Schema](
     tableName: String,
     key: PrimaryKey,
     projections: ProjectionExpression[_, _]*
@@ -515,13 +496,24 @@ object DynamoDBQuery {
       UpdateExpression(action)
     )
 
-  def update[A: Schema](tableName: String, key: PrimaryKey)(action: Action[A]): DynamoDBQuery[A, Option[A]] =
+  private[dynamodb] def update[A: Schema](tableName: String, key: PrimaryKey)(
+    action: Action[A]
+  ): DynamoDBQuery[A, Option[A]] =
     updateItem(tableName, key)(action).map(_.flatMap(item => fromItem(item).toOption))
 
+  def update[From: Schema, Pk, Sk](tableName: String)(primaryKeyExpr: KeyConditionExpr.PrimaryKeyExpr[From, Pk, Sk])(
+    action: Action[From]
+  ): DynamoDBQuery[From, Option[From]] =
+    updateItem(tableName, primaryKeyExpr.asAttrMap)(action).map(_.flatMap(item => fromItem(item).toOption))
+
   def deleteItem(tableName: String, key: PrimaryKey): Write[Any, Option[Item]] = DeleteItem(TableName(tableName), key)
 
-  def delete[A: Schema](tableName: String, key: PrimaryKey): DynamoDBQuery[Any, Option[A]] =
-    deleteItem(tableName, key).map(_.flatMap(item => fromItem(item).toOption))
+  def delete[From: Schema, Pk, Sk](
+    tableName: String
+  )(
+    primaryKeyExpr: KeyConditionExpr.PrimaryKeyExpr[From, Pk, Sk]
+  ): DynamoDBQuery[Any, Option[From]] =
+    deleteItem(tableName, primaryKeyExpr.asAttrMap).map(_.flatMap(item => fromItem(item).toOption))
 
   /**
    * when executed will return a Tuple of {{{(Chunk[Item], LastEvaluatedKey)}}}
@@ -846,7 +838,7 @@ object DynamoDBQuery {
     exclusiveStartKey: LastEvaluatedKey =
       None,                                                     // allows client to control start position - eg for client managed paging
     filterExpression: Option[FilterExpression[_]] = None,
-    keyConditionExpression: Option[KeyConditionExpression] = None,
+    keyConditionExpr: Option[KeyConditionExpr[_, _, _]] = None,
     projections: List[ProjectionExpression[_, _]] = List.empty, // if empty all attributes will be returned
     capacity: ReturnConsumedCapacity = ReturnConsumedCapacity.None,
     select: Option[Select] = None,                              // if ProjectExpression supplied then only valid value is SpecificAttributes
@@ -879,7 +871,7 @@ object DynamoDBQuery {
     exclusiveStartKey: LastEvaluatedKey =
       None,                                                     // allows client to control start position - eg for client managed paging
     filterExpression: Option[FilterExpression[_]] = None,
-    keyConditionExpression: Option[KeyConditionExpression] = None,
+    keyConditionExpr: Option[KeyConditionExpr[_, _, _]] = None,
     projections: List[ProjectionExpression[_, _]] = List.empty, // if empty all attributes will be returned
     capacity: ReturnConsumedCapacity = ReturnConsumedCapacity.None,
     select: Option[Select] = None,                              // if ProjectExpression supplied then only valid value is SpecificAttributes

dynamodb/src/main/scala/zio/dynamodb/KeyConditionExpr.scala

@@ -0,0 +1,147 @@
+package zio.dynamodb
+
+/**
+ * This sum type models:
+ * 1) partition key equality expressions
+ * 2) composite primary key expressions where sort key expression is equality
+ * 3) extended composite primary key expressions where sort key is not equality eg >, <, >=, <=, between, begins_with
+ *
+ * Note 1), 2) and 3) are all valid key condition expressions used in Query DynamoDB queries
+ * BUT only 1) and 2) are valid primary key expressions that can be used in GetItem, UpdateItem and DeleteItem DynamoDB queries
+ */
+sealed trait KeyConditionExpr[-From, +Pk, +Sk] extends Renderable { self =>
+  def render: AliasMapRender[String]
+}
+
+object KeyConditionExpr {
+  type SortKeyNotUsed
+
+  sealed trait PrimaryKeyExpr[-From, +Pk, +Sk] extends KeyConditionExpr[From, Pk, Sk] {
+    def asAttrMap: AttrMap
+  }
+
+  def getOrInsert[From, To](primaryKeyName: String): AliasMapRender[String] =
+    // note primary keys must be scalar values, they can't be nested
+    AliasMapRender.getOrInsert(ProjectionExpression.MapElement[From, To](ProjectionExpression.Root, primaryKeyName))
+
+  private[dynamodb] final case class PartitionKeyEquals[-From, +Pk](pk: PartitionKey[From, Pk], value: AttributeValue)
+      extends PrimaryKeyExpr[From, Pk, SortKeyNotUsed] { self =>
+
+    def &&[From1 <: From, Sk](other: SortKeyEquals[From1, Sk]): CompositePrimaryKeyExpr[From1, Pk, Sk] =
+      CompositePrimaryKeyExpr[From1, Pk, Sk](self.asInstanceOf[PartitionKeyEquals[From1, Pk]], other)
+    def &&[From1 <: From, Sk](
+      other: ExtendedSortKeyExpr[From1, Sk]
+    ): ExtendedCompositePrimaryKeyExpr[From1, Pk, Sk]                                                  =
+      ExtendedCompositePrimaryKeyExpr[From1, Pk, Sk](self.asInstanceOf[PartitionKeyEquals[From1, Pk]], other)
+
+    def asAttrMap: AttrMap = AttrMap(pk.keyName -> value)
+
+    override def render: AliasMapRender[String] =
+      for {
+        v        <- AliasMapRender.getOrInsert(value)
+        keyAlias <- KeyConditionExpr.getOrInsert(pk.keyName)
+      } yield s"${keyAlias} = $v"
+
+  }
+
+  private[dynamodb] final case class SortKeyEquals[-From, +Sk](sortKey: SortKey[From, Sk], value: AttributeValue) {
+    self =>
+    def miniRender: AliasMapRender[String] =
+      for {
+        v        <- AliasMapRender.getOrInsert(value)
+        keyAlias <- KeyConditionExpr.getOrInsert(sortKey.keyName)
+      } yield s"${keyAlias} = $v"
+  }
+
+  private[dynamodb] final case class CompositePrimaryKeyExpr[-From, +Pk, +Sk](
+    pk: PartitionKeyEquals[From, Pk],
+    sk: SortKeyEquals[From, Sk]
+  ) extends PrimaryKeyExpr[From, Pk, Sk] {
+    self =>
+
+    def asAttrMap: AttrMap = PrimaryKey(pk.pk.keyName -> pk.value, sk.sortKey.keyName -> sk.value)
+
+    override def render: AliasMapRender[String] =
+      for {
+        pkStr <- pk.render
+        skStr <- sk.miniRender
+      } yield s"$pkStr AND $skStr"
+
+  }
+  private[dynamodb] final case class ExtendedCompositePrimaryKeyExpr[-From, +Pk, +Sk](
+    pk: PartitionKeyEquals[From, Pk],
+    sk: ExtendedSortKeyExpr[From, Sk]
+  ) extends KeyConditionExpr[From, Pk, Sk] {
+    self =>
+
+    def render: AliasMapRender[String] =
+      for {
+        pkStr <- pk.render
+        skStr <- sk.miniRender
+      } yield s"$pkStr AND $skStr"
+
+  }
+
+  sealed trait ExtendedSortKeyExpr[-From, +Sk] { self =>
+    def miniRender: AliasMapRender[String] =
+      self match {
+        case ExtendedSortKeyExpr.GreaterThan(sk, value)        =>
+          for {
+            v        <- AliasMapRender.getOrInsert(value)
+            keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
+          } yield s"${keyAlias} > $v"
+        case ExtendedSortKeyExpr.LessThan(sk, value)           =>
+          for {
+            v        <- AliasMapRender.getOrInsert(value)
+            keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
+          } yield s"${keyAlias} < $v"
+        case ExtendedSortKeyExpr.NotEqual(sk, value)           =>
+          for {
+            v        <- AliasMapRender.getOrInsert(value)
+            keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
+          } yield s"${keyAlias} <> $v"
+        case ExtendedSortKeyExpr.LessThanOrEqual(sk, value)    =>
+          for {
+            v        <- AliasMapRender.getOrInsert(value)
+            keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
+          } yield s"${keyAlias} <= $v"
+        case ExtendedSortKeyExpr.GreaterThanOrEqual(sk, value) =>
+          for {
+            v        <- AliasMapRender.getOrInsert(value)
+            keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
+          } yield s"${keyAlias} >= $v"
+        case ExtendedSortKeyExpr.Between(left, min, max)       =>
+          for {
+            min2     <- AliasMapRender.getOrInsert(min)
+            max2     <- AliasMapRender.getOrInsert(max)
+            keyAlias <- KeyConditionExpr.getOrInsert(left.keyName)
+          } yield s"${keyAlias} BETWEEN $min2 AND $max2"
+        case ExtendedSortKeyExpr.BeginsWith(left, value)       =>
+          for {
+            v        <- AliasMapRender.getOrInsert(value)
+            keyAlias <- KeyConditionExpr.getOrInsert(left.keyName)
+          } yield s"begins_with(${keyAlias}, $v)"
+      }
+
+  }
+  object ExtendedSortKeyExpr {
+    private[dynamodb] final case class GreaterThan[From, +To](sortKey: SortKey[From, To], value: AttributeValue)
+        extends ExtendedSortKeyExpr[From, To]
+    private[dynamodb] final case class LessThan[From, +To](sortKey: SortKey[From, To], value: AttributeValue)
+        extends ExtendedSortKeyExpr[From, To]
+    private[dynamodb] final case class NotEqual[From, +To](sortKey: SortKey[From, To], value: AttributeValue)
+        extends ExtendedSortKeyExpr[From, To]
+    private[dynamodb] final case class LessThanOrEqual[From, +To](sortKey: SortKey[From, To], value: AttributeValue)
+        extends ExtendedSortKeyExpr[From, To]
+    private[dynamodb] final case class GreaterThanOrEqual[From, +To](sortKey: SortKey[From, To], value: AttributeValue)
+        extends ExtendedSortKeyExpr[From, To]
+    private[dynamodb] final case class Between[From, +To](
+      left: SortKey[From, To],
+      min: AttributeValue,
+      max: AttributeValue
+    ) extends ExtendedSortKeyExpr[From, To]
+    private[dynamodb] final case class BeginsWith[From, +To](left: SortKey[From, To], value: AttributeValue)
+        extends ExtendedSortKeyExpr[From, To]
+  }
+
+}