Skip to content

Commit

Permalink
fix KeyConditionExpression rendering for names
Browse files Browse the repository at this point in the history
  • Loading branch information
@googley42
googley42 committed Jul 26, 2023
1 parent f5a63d3 commit 185da86
Show file tree
Hide file tree
Showing 3 changed files with 176 additions and 79 deletions.
78 changes: 78 additions & 0 deletions dynamodb/src/it/scala/zio/dynamodb/LiveSpec.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import software.amazon.awssdk.services.dynamodb.model.{ DynamoDbException, Idemp
import zio.dynamodb.UpdateExpression.Action.SetAction
import zio.dynamodb.UpdateExpression.SetOperand
import zio.dynamodb.PartitionKeyExpression.PartitionKey
import zio.dynamodb.KeyConditionExpression.partitionKey
import zio.dynamodb.KeyConditionExpression.sortKey
import zio.dynamodb.SortKeyExpression.SortKey
import zio.aws.{ dynamodb, netty }
import zio._
Expand Down Expand Up @@ -116,6 +118,12 @@ object LiveSpec extends ZIOSpecDefault {
AttributeDefinition.attrDefnString(name)
)

def sortKeyStringTableWithKeywords(tableName: String) =
createTable(tableName, KeySchema("and", "source"), BillingMode.PayPerRequest)(
AttributeDefinition.attrDefnString("and"),
AttributeDefinition.attrDefnString("source")
)

private def managedTable(tableDefinition: String => CreateTable) =
ZIO
.acquireRelease(
Expand Down Expand Up @@ -150,6 +158,17 @@ object LiveSpec extends ZIOSpecDefault {
}
}

def withKeywordsTable(
f: String => ZIO[DynamoDBExecutor, Throwable, TestResult]
) =
ZIO.scoped {
managedTable(sortKeyStringTableWithKeywords).flatMap { table =>
for {
result <- f(table.value)
} yield result
}
}

def withDefaultAndNumberTables(
f: (String, String) => ZIO[DynamoDBExecutor, Throwable, TestResult]
) =
Expand Down Expand Up @@ -179,8 +198,67 @@ object LiveSpec extends ZIOSpecDefault {
val (id, num, ttl) = ProjectionExpression.accessors[ExpressionAttrNames]
}

final case class ExpressionAttrNames2(and: String, source: String, ttl: Option[Long])
object ExpressionAttrNames2 {
implicit val schema: Schema.CaseClass3[String, String, Option[Long], ExpressionAttrNames2] =
DeriveSchema.gen[ExpressionAttrNames2]
val (and, source, ttl) = ProjectionExpression.accessors[ExpressionAttrNames2]
}

val debugSuite = suite("debug")(
test("queryAll should handle keywords in primary key names using high level API") {
withKeywordsTable { tableName =>
val query = DynamoDBQuery
.queryAll[ExpressionAttrNames2](tableName)
.whereKey(ExpressionAttrNames2.and === "and1" && ExpressionAttrNames2.source === "source1")
.filter(ExpressionAttrNames2.ttl.notExists)
query.execute.flatMap(_.runDrain).exit.map { result =>
assert(result)(succeeds(isUnit))
}
}
},
test("queryAll should handle keywords in primary key names using low level API") {
withKeywordsTable { tableName =>
val query = DynamoDBQuery
.queryAll[ExpressionAttrNames2](tableName)
.whereKey(partitionKey("and") === "and1" && sortKey("source") === "source1")
.filter(ExpressionAttrNames2.ttl.notExists)
query.execute.flatMap(_.runDrain).exit.map { result =>
assert(result)(succeeds(isUnit))
}
}
}
)
.provideSomeLayerShared[TestEnvironment](
testLayer.orDie
) @@ nondeterministic

val mainSuite: Spec[TestEnvironment, Any] =
suite("live test")(
suite("key words in Key Condition Expressions")(
test("queryAll should handle keywords in primary key name using high level API") {
withKeywordsTable { tableName =>
val query = DynamoDBQuery
.queryAll[ExpressionAttrNames2](tableName)
.whereKey(ExpressionAttrNames2.and === "and1" && ExpressionAttrNames2.source === "source1")
.filter(ExpressionAttrNames2.ttl.notExists)
query.execute.flatMap(_.runDrain).exit.map { result =>
assert(result)(succeeds(isUnit))
}
}
},
test("queryAll should handle keywords in primary key name using low level API") {
withKeywordsTable { tableName =>
val query = DynamoDBQuery
.queryAll[ExpressionAttrNames2](tableName)
.whereKey(partitionKey("and") === "and1" && sortKey("source") === "source1")
.filter(ExpressionAttrNames2.ttl.notExists)
query.execute.flatMap(_.runDrain).exit.map { result =>
assert(result)(succeeds(isUnit))
}
}
}
),
suite("keywords in expression attribute names")(
suite("using high level api")(
test("scanAll should handle keyword") {
Expand Down
86 changes: 42 additions & 44 deletions dynamodb/src/main/scala/zio/dynamodb/KeyConditionExpression.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,13 @@ sealed trait KeyConditionExpression extends Renderable { self =>
}

object KeyConditionExpression {

def getOrInsert[From](primaryKeyName: String): AliasMapRender[String] =
AliasMapRender.getOrInsert(ProjectionExpression.MapElement[From, String](Root, primaryKeyName))
private[dynamodb] final case class And(left: PartitionKeyExpression, right: SortKeyExpression)
extends KeyConditionExpression
def partitionKey(key: String): PartitionKey = PartitionKey(key)
def partitionKey(key: String): PartitionKey = PartitionKey(key)
def sortKey(key: String): SortKey = SortKey(key)

/**
* Create a KeyConditionExpression from a ConditionExpression
Expand Down Expand Up @@ -156,7 +160,10 @@ sealed trait PartitionKeyExpression extends KeyConditionExpression { self =>
override def render: AliasMapRender[String] =
self match {
case PartitionKeyExpression.Equals(left, right) =>
AliasMapRender.getOrInsert(right).map(v => s"${left.keyName} = $v")
for {
v <- AliasMapRender.getOrInsert(right)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} = $v"
}
}
object PartitionKeyExpression {
Expand All @@ -171,55 +178,46 @@ sealed trait SortKeyExpression { self =>
def render: AliasMapRender[String] =
self match {
case SortKeyExpression.Equals(left, right) =>
AliasMapRender
.getOrInsert(right)
.map { v =>
s"${left.keyName} = $v"
}
for {
v <- AliasMapRender.getOrInsert(right)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} = $v"
case SortKeyExpression.LessThan(left, right) =>
AliasMapRender
.getOrInsert(right)
.map { v =>
s"${left.keyName} < $v"
}
for {
v <- AliasMapRender.getOrInsert(right)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} < $v"
case SortKeyExpression.NotEqual(left, right) =>
AliasMapRender
.getOrInsert(right)
.map { v =>
s"${left.keyName} <> $v"
}
for {
v <- AliasMapRender.getOrInsert(right)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} <> $v"
case SortKeyExpression.GreaterThan(left, right) =>
AliasMapRender
.getOrInsert(right)
.map { v =>
s"${left.keyName} > $v"
}
for {
v <- AliasMapRender.getOrInsert(right)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} > $v"
case SortKeyExpression.LessThanOrEqual(left, right) =>
AliasMapRender
.getOrInsert(right)
.map { v =>
s"${left.keyName} <= $v"
}
for {
v <- AliasMapRender.getOrInsert(right)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} <= $v"
case SortKeyExpression.GreaterThanOrEqual(left, right) =>
AliasMapRender
.getOrInsert(right)
.map { v =>
s"${left.keyName} >= $v"
}
for {
v <- AliasMapRender.getOrInsert(right)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} >= $v"
case SortKeyExpression.Between(left, min, max) =>
AliasMapRender
.getOrInsert(min)
.flatMap(min =>
AliasMapRender.getOrInsert(max).map { max =>
s"${left.keyName} BETWEEN $min AND $max"
}
)
for {
min2 <- AliasMapRender.getOrInsert(min)
max2 <- AliasMapRender.getOrInsert(max)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"${keyName} BETWEEN $min2 AND $max2"
case SortKeyExpression.BeginsWith(left, value) =>
AliasMapRender
.getOrInsert(value)
.map { v =>
s"begins_with(${left.keyName}, $v)"
}
for {
v <- AliasMapRender.getOrInsert(value)
keyName <- KeyConditionExpression.getOrInsert(left.keyName)
} yield s"begins_with(${keyName}, $v)"
}

}
Expand Down
Loading

0 comments on commit 185da86

Please sign in to comment.