Skip to content

Commit

Permalink
port over key condition render fix from main
Browse files Browse the repository at this point in the history
  • Loading branch information
@googley42
googley42 committed Jul 29, 2023
1 parent fba3c02 commit e4646b3
Show file tree
Hide file tree
Showing 3 changed files with 171 additions and 65 deletions.
74 changes: 74 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 @@ -114,6 +114,12 @@ object LiveSpec extends ZIOSpecDefault {
AttributeDefinition.attrDefnString(name)
)

private 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 All @@ -123,6 +129,17 @@ object LiveSpec extends ZIOSpecDefault {
} yield TableName(tableName)
)(tName => deleteTable(tName.value).execute.orDie)

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

// TODO: Avi - fix problem with inference of this function when splitting suites
// "a type was inferred to be `Any`; this may indicate a programming error."
private def withTemporaryTable[R](
Expand Down Expand Up @@ -177,6 +194,13 @@ object LiveSpec extends ZIOSpecDefault {
val (id, num, ttl) = ProjectionExpression.accessors[ExpressionAttrNames]
}

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

val debugSuite = suite("debug")(
test("delete should handle keyword") {
withDefaultTable { tableName =>
Expand All @@ -198,6 +222,56 @@ object LiveSpec extends ZIOSpecDefault {

val mainSuite: Spec[TestEnvironment, Any] =
suite("live test")(
suite("key words in Key Condition Expressions")(
test("queryAll should handle keywords in primary key names using high level API") {
withPkKeywordsTable { tableName =>
val query = DynamoDBQuery
.queryAll[ExpressionAttrNamesPkKeywords](tableName)
.whereKey(
ExpressionAttrNamesPkKeywords.and.partitionKey === "and1" && ExpressionAttrNamesPkKeywords.source.sortKey === "source1"
)
.filter(ExpressionAttrNamesPkKeywords.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") {
withPkKeywordsTable { tableName =>
val query = DynamoDBQuery
.queryAll[ExpressionAttrNamesPkKeywords](tableName)
.whereKey($("and").partitionKey === "and1" && $("source").sortKey === "source1")
.filter(ExpressionAttrNamesPkKeywords.ttl.notExists)
query.execute.flatMap(_.runDrain).exit.map { result =>
assert(result)(succeeds(isUnit))
}
}
},
test("querySome should handle keywords in primary key name using high level API") {
withPkKeywordsTable { tableName =>
val query = DynamoDBQuery
.querySome[ExpressionAttrNamesPkKeywords](tableName, 1)
.whereKey(
ExpressionAttrNamesPkKeywords.and.partitionKey === "and1" && ExpressionAttrNamesPkKeywords.source.sortKey === "source1"
)
.filter(ExpressionAttrNamesPkKeywords.ttl.notExists)
for {
result <- query.execute
} yield assert(result._1)(hasSize(equalTo(0)))
}
},
test("querySome should handle keywords in primary key name using low level API") {
withPkKeywordsTable { tableName =>
val query = DynamoDBQuery
.querySome[ExpressionAttrNames](tableName, 1)
.whereKey($("and").partitionKey === "and1" && $("source").sortKey === "source1")
.filter(ExpressionAttrNames.ttl.notExists)
for {
result <- query.execute
} yield assert(result._1)(hasSize(equalTo(0)))
}
}
),
suite("keywords in expression attribute names")(
suite("using high level api")(
test("scanAll should handle keyword") {
Expand Down
72 changes: 41 additions & 31 deletions dynamodb/src/main/scala/zio/dynamodb/KeyConditionExpr.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,9 @@ sealed trait KeyConditionExpr[-From, +To] extends Renderable { self =>
}

object KeyConditionExpr {
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, +To](pk: PartitionKey[From, To], value: AttributeValue)
extends KeyConditionExpr[From, To] { self =>
Expand All @@ -26,15 +29,20 @@ object KeyConditionExpr {
def asAttrMap: AttrMap = AttrMap(pk.keyName -> value)

override def render: AliasMapRender[String] =
AliasMapRender.getOrInsert(value).map(v => s"${pk.keyName} = $v")
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(pk.keyName)
} yield s"${keyAlias} = $v"

}

private[dynamodb] final case class SortKeyEquals[-From, +To](sortKey: SortKey[From, To], value: AttributeValue) {
self =>
def miniRender: AliasMapRender[String] =
AliasMapRender
.getOrInsert(value)
.map(v => s"${sortKey.keyName} = $v")
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(sortKey.keyName)
} yield s"${keyAlias} = $v"
}

private[dynamodb] final case class CompositePrimaryKeyExpr[-From](
Expand Down Expand Up @@ -70,39 +78,41 @@ object KeyConditionExpr {
def miniRender: AliasMapRender[String] =
self match {
case ExtendedSortKeyExpr.GreaterThan(sk, value) =>
AliasMapRender
.getOrInsert(value)
.map(v => s"${sk.keyName} > $v")
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
} yield s"${keyAlias} > $v"
case ExtendedSortKeyExpr.LessThan(sk, value) =>
AliasMapRender
.getOrInsert(value)
.map(v => s"${sk.keyName} < $v")
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
} yield s"${keyAlias} < $v"
case ExtendedSortKeyExpr.NotEqual(sk, value) =>
AliasMapRender
.getOrInsert(value)
.map(v => s"${sk.keyName} <> $v")
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
} yield s"${keyAlias} <> $v"
case ExtendedSortKeyExpr.LessThanOrEqual(sk, value) =>
AliasMapRender
.getOrInsert(value)
.map(v => s"${sk.keyName} <= $v")
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
} yield s"${keyAlias} <= $v"
case ExtendedSortKeyExpr.GreaterThanOrEqual(sk, value) =>
AliasMapRender
.getOrInsert(value)
.map(v => s"${sk.keyName} >= $v")
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(sk.keyName)
} yield s"${keyAlias} >= $v"
case ExtendedSortKeyExpr.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)
keyAlias <- KeyConditionExpr.getOrInsert(left.keyName)
} yield s"${keyAlias} BETWEEN $min2 AND $max2"
case ExtendedSortKeyExpr.BeginsWith(left, value) =>
AliasMapRender
.getOrInsert(value)
.map { v =>
s"begins_with(${left.keyName}, $v)"
}
for {
v <- AliasMapRender.getOrInsert(value)
keyAlias <- KeyConditionExpr.getOrInsert(left.keyName)
} yield s"begins_with(${keyAlias}, $v)"
}

}
Expand Down
90 changes: 56 additions & 34 deletions dynamodb/src/test/scala/zio/dynamodb/AliasMapRenderSpec.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -323,126 +323,148 @@ object AliasMapRenderSpec extends ZIOSpecDefault {
suite("Sort key expressions")(
test("Equals") {
val map = Map(
avKey(one) -> ":v0"
avKey(one) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) =
KeyConditionExpr.SortKeyEquals($("num").sortKey, one).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("num = :v0"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("#n1 = :v0"))
},
test("LessThan") {
val map = Map(
avKey(one) -> ":v0"
avKey(one) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) =
KeyConditionExpr.ExtendedSortKeyExpr.LessThan($("num").sortKey, one).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("num < :v0"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("#n1 < :v0"))
},
test("NotEqual") {
val map = Map(
avKey(one) -> ":v0"
avKey(one) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) =
KeyConditionExpr.ExtendedSortKeyExpr.NotEqual($("num").sortKey, one).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("num <> :v0"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("#n1 <> :v0"))
},
test("GreaterThan") {
val map = Map(
avKey(one) -> ":v0"
avKey(one) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) =
KeyConditionExpr.ExtendedSortKeyExpr.GreaterThan($("num").sortKey, one).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("num > :v0"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("#n1 > :v0"))
},
test("LessThanOrEqual") {
val map = Map(
avKey(one) -> ":v0"
avKey(one) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) =
KeyConditionExpr.ExtendedSortKeyExpr.LessThanOrEqual($("num").sortKey, one).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("num <= :v0"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("#n1 <= :v0"))
},
test("GreaterThanOrEqual") {
val map = Map(
avKey(one) -> ":v0"
avKey(one) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) =
KeyConditionExpr.ExtendedSortKeyExpr.GreaterThanOrEqual($("num").sortKey, one).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("num >= :v0"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("#n1 >= :v0"))
},
test("Between") {
val map = Map(
avKey(one) -> ":v0",
avKey(two) -> ":v1"
avKey(one) -> ":v0",
avKey(two) -> ":v1",
pathSegment(Root, "num") -> "#n2",
fullPath($("num")) -> "#n2"
)

val (aliasMap, expression) =
KeyConditionExpr.ExtendedSortKeyExpr.Between($("num").sortKey, one, two).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("num BETWEEN :v0 AND :v1"))
assert(aliasMap)(equalTo(AliasMap(map, 3))) &&
assert(expression)(equalTo("#n2 BETWEEN :v0 AND :v1"))
},
test("BeginsWith") {
val map = Map(
avKey(name) -> ":v0"
avKey(name) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) =
KeyConditionExpr.ExtendedSortKeyExpr.BeginsWith($("num").sortKey, name).miniRender.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("begins_with(num, :v0)"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("begins_with(#n1, :v0)"))
}
),
suite("PartitionKeyExpression")(
test("Equals") {
val map = Map(
avKey(one) -> ":v0"
avKey(one) -> ":v0",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1"
)

val (aliasMap, expression) = KeyConditionExpr
.PartitionKeyEquals($("num").partitionKey, one)
.render
.execute

assert(aliasMap)(equalTo(AliasMap(map, 1))) &&
assert(expression)(equalTo("num = :v0"))
assert(aliasMap)(equalTo(AliasMap(map, 2))) &&
assert(expression)(equalTo("#n1 = :v0"))
}
),
test("And") {
val map = Map(
avKey(two) -> ":v0",
avKey(one) -> ":v1",
avKey(three) -> ":v2"
avKey(two) -> ":v0",
avKey(one) -> ":v2",
avKey(three) -> ":v3",
pathSegment(Root, "num") -> "#n1",
fullPath($("num")) -> "#n1",
pathSegment(Root, "num2") -> "#n4",
fullPath($("num2")) -> "#n4"
)

val (aliasMap, expression) = KeyConditionExpr
.ExtendedCompositePrimaryKeyExpr(
KeyConditionExpr.PartitionKeyEquals($("num").partitionKey, two),
KeyConditionExpr.ExtendedSortKeyExpr.Between($("num").sortKey, one, three)
KeyConditionExpr.ExtendedSortKeyExpr.Between($("num2").sortKey, one, three)
)
.render
.execute

assert(aliasMap)(equalTo(AliasMap(map, 3))) &&
assert(expression)(equalTo("num = :v0 AND num BETWEEN :v1 AND :v2"))
assert(aliasMap)(equalTo(AliasMap(map, 5))) &&
assert(expression)(equalTo("#n1 = :v0 AND #n4 BETWEEN :v2 AND :v3"))
}
),
suite("AttributeValueType")(
Expand Down

0 comments on commit e4646b3

Please sign in to comment.