Skip to content

Commit 5f75f6f

Browse files
tachyonicstachyonics
authored
Merge pull request #8 from amzn/gsi_query
Provide the ability to query an index name and on Keys Only projections.
2 parents 73f2154 + afb3017 commit 5f75f6f

17 files changed

+647
-84
lines changed

‎Package.resolved

Lines changed: 13 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Sources/SmokeDynamoDB/AWSDynamoDBKeysProjection+DynamoDBKeysProjectionAsync.swift

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License").
4+
// You may not use this file except in compliance with the License.
5+
// A copy of the License is located at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// or in the "license" file accompanying this file. This file is distributed
10+
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
11+
// express or implied. See the License for the specific language governing
12+
// permissions and limitations under the License.
13+
//
14+
// AWSDynamoDBKeysProjection+DynamoDBKeysProjectionAsync.swift
15+
// SmokeDynamoDB
16+
//
17+
18+
import Foundation
19+
import SmokeAWSCore
20+
import DynamoDBModel
21+
import SmokeHTTPClient
22+
import LoggerAPI
23+
24+
/// DynamoDBKeysProjection conformance async functions
25+
public extension AWSDynamoDBKeysProjection {
26+
27+
public func queryAsync<AttributesType>(
28+
forPartitionKey partitionKey: String,
29+
sortKeyCondition: AttributeCondition?,
30+
completion: @escaping (HTTPResult<[CompositePrimaryKey<AttributesType>]>) -> ())
31+
throws where AttributesType: PrimaryKeyAttributes {
32+
let partialResults = QueryPaginationResults<AttributesType>()
33+
34+
try partialQueryAsync(forPartitionKey: partitionKey,
35+
sortKeyCondition: sortKeyCondition,
36+
partialResults: partialResults,
37+
completion: completion)
38+
}
39+
40+
private func partialQueryAsync<AttributesType>(
41+
forPartitionKey partitionKey: String,
42+
sortKeyCondition: AttributeCondition?,
43+
partialResults: QueryPaginationResults<AttributesType>,
44+
completion: @escaping (HTTPResult<[CompositePrimaryKey<AttributesType>]>) -> ())
45+
throws where AttributesType: PrimaryKeyAttributes {
46+
func handleQueryResult(result: HTTPResult<([CompositePrimaryKey<AttributesType>], String?)>) {
47+
switch result {
48+
case .response(let paginatedItems):
49+
partialResults.items += paginatedItems.0
50+
51+
// if there are more items
52+
if let lastEvaluatedKey = paginatedItems.1 {
53+
partialResults.exclusiveStartKey = lastEvaluatedKey
54+
55+
do {
56+
try partialQueryAsync(forPartitionKey: partitionKey,
57+
sortKeyCondition: sortKeyCondition,
58+
partialResults: partialResults,
59+
completion: completion)
60+
} catch {
61+
completion(.error(error))
62+
}
63+
} else {
64+
// we have all the items
65+
completion(.response(partialResults.items))
66+
}
67+
case .error(let error):
68+
completion(.error(error))
69+
}
70+
}
71+
72+
try queryAsync(forPartitionKey: partitionKey,
73+
sortKeyCondition: sortKeyCondition,
74+
limit: nil,
75+
exclusiveStartKey: partialResults.exclusiveStartKey,
76+
completion: handleQueryResult)
77+
}
78+
79+
public func queryAsync<AttributesType>(
80+
forPartitionKey partitionKey: String,
81+
sortKeyCondition: AttributeCondition?,
82+
limit: Int?,
83+
exclusiveStartKey: String?,
84+
completion: @escaping (HTTPResult<([CompositePrimaryKey<AttributesType>], String?)>) -> ())
85+
throws where AttributesType: PrimaryKeyAttributes {
86+
let queryInput = try DynamoDBModel.QueryInput.forSortKeyCondition(forPartitionKey: partitionKey, targetTableName: targetTableName,
87+
primaryKeyType: AttributesType.self,
88+
sortKeyCondition: sortKeyCondition, limit: limit,
89+
exclusiveStartKey: exclusiveStartKey)
90+
try dynamodb.queryAsync(input: queryInput) { result in
91+
switch result {
92+
case .response(let queryOutput):
93+
let lastEvaluatedKey: String?
94+
if let returnedLastEvaluatedKey = queryOutput.lastEvaluatedKey {
95+
let encodedLastEvaluatedKey: Data
96+
97+
do {
98+
encodedLastEvaluatedKey = try AWSDynamoDBTable.jsonEncoder.encode(returnedLastEvaluatedKey)
99+
} catch {
100+
return completion(.error(error))
101+
}
102+
103+
lastEvaluatedKey = String(data: encodedLastEvaluatedKey, encoding: .utf8)
104+
} else {
105+
lastEvaluatedKey = nil
106+
}
107+
108+
if let outputAttributeValues = queryOutput.items {
109+
let items: [CompositePrimaryKey<AttributesType>]
110+
111+
do {
112+
items = try outputAttributeValues.map { values in
113+
let attributeValue = DynamoDBModel.AttributeValue(M: values)
114+
115+
return try AWSDynamoDBTable.dynamodbDecoder.decode(attributeValue)
116+
}
117+
} catch {
118+
return completion(.error(error))
119+
}
120+
121+
completion(.response((items, lastEvaluatedKey)))
122+
} else {
123+
completion(.response(([], lastEvaluatedKey)))
124+
}
125+
case .error(let error):
126+
return completion(.error(error))
127+
}
128+
}
129+
}
130+
}

‎Sources/SmokeDynamoDB/AWSDynamoDBKeysProjection+DynamoDBKeysProjectionSync.swift

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License").
4+
// You may not use this file except in compliance with the License.
5+
// A copy of the License is located at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// or in the "license" file accompanying this file. This file is distributed
10+
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
11+
// express or implied. See the License for the specific language governing
12+
// permissions and limitations under the License.
13+
//
14+
// AWSDynamoDBKeysProjection+DynamoDBKeysProjectionSync.swift
15+
// SmokeDynamoDB
16+
//
17+
18+
import Foundation
19+
import SmokeAWSCore
20+
import DynamoDBModel
21+
import SmokeHTTPClient
22+
import LoggerAPI
23+
24+
/// DynamoDBKeysProjection conformance sync functions
25+
public extension AWSDynamoDBKeysProjection {
26+
27+
public func querySync<AttributesType>(forPartitionKey partitionKey: String,
28+
sortKeyCondition: AttributeCondition?) throws
29+
-> [CompositePrimaryKey<AttributesType>] {
30+
31+
var items: [CompositePrimaryKey<AttributesType>] = []
32+
var exclusiveStartKey: String?
33+
34+
while true {
35+
let paginatedItems: ([CompositePrimaryKey<AttributesType>], String?) =
36+
try querySync(forPartitionKey: partitionKey,
37+
sortKeyCondition: sortKeyCondition,
38+
limit: nil,
39+
exclusiveStartKey: exclusiveStartKey)
40+
41+
items += paginatedItems.0
42+
43+
// if there are more items
44+
if let lastEvaluatedKey = paginatedItems.1 {
45+
exclusiveStartKey = lastEvaluatedKey
46+
} else {
47+
// we have all the items
48+
return items
49+
}
50+
}
51+
}
52+
53+
public func querySync<AttributesType>(forPartitionKey partitionKey: String,
54+
sortKeyCondition: AttributeCondition?,
55+
limit: Int?,
56+
exclusiveStartKey: String?) throws
57+
-> ([CompositePrimaryKey<AttributesType>], String?)
58+
where AttributesType: PrimaryKeyAttributes {
59+
let queryInput = try DynamoDBModel.QueryInput.forSortKeyCondition(forPartitionKey: partitionKey, targetTableName: targetTableName,
60+
primaryKeyType: AttributesType.self,
61+
sortKeyCondition: sortKeyCondition, limit: limit,
62+
exclusiveStartKey: exclusiveStartKey)
63+
let queryOutput = try dynamodb.querySync(input: queryInput)
64+
65+
let lastEvaluatedKey: String?
66+
if let returnedLastEvaluatedKey = queryOutput.lastEvaluatedKey {
67+
let encodedLastEvaluatedKey = try AWSDynamoDBTable.jsonEncoder.encode(returnedLastEvaluatedKey)
68+
69+
lastEvaluatedKey = String(data: encodedLastEvaluatedKey, encoding: .utf8)
70+
} else {
71+
lastEvaluatedKey = nil
72+
}
73+
74+
if let outputAttributeValues = queryOutput.items {
75+
let items: [CompositePrimaryKey<AttributesType>] = try outputAttributeValues.map { values in
76+
let attributeValue = DynamoDBModel.AttributeValue(M: values)
77+
78+
return try AWSDynamoDBTable.dynamodbDecoder.decode(attributeValue)
79+
}
80+
81+
return (items, lastEvaluatedKey)
82+
} else {
83+
return ([], lastEvaluatedKey)
84+
}
85+
}
86+
}

‎Sources/SmokeDynamoDB/AWSDynamoDBKeysProjection.swift

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License").
4+
// You may not use this file except in compliance with the License.
5+
// A copy of the License is located at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// or in the "license" file accompanying this file. This file is distributed
10+
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
11+
// express or implied. See the License for the specific language governing
12+
// permissions and limitations under the License.
13+
//
14+
// AWSDynamoDBKeysProjection.swift
15+
// SmokeDynamoDB
16+
//
17+
18+
import Foundation
19+
import LoggerAPI
20+
import DynamoDBClient
21+
import DynamoDBModel
22+
import SmokeAWSCore
23+
import SmokeHTTPClient
24+
25+
public class AWSDynamoDBKeysProjection: DynamoDBKeysProjection {
26+
internal let dynamodb: AWSDynamoDBClient
27+
internal let targetTableName: String
28+
29+
static internal let dynamodbEncoder = DynamoDBEncoder()
30+
static internal let dynamodbDecoder = DynamoDBDecoder()
31+
static internal let jsonEncoder = JSONEncoder()
32+
static internal let jsonDecoder = JSONDecoder()
33+
34+
internal class QueryPaginationResults<AttributesType: PrimaryKeyAttributes> {
35+
var items: [CompositePrimaryKey<AttributesType>] = []
36+
var exclusiveStartKey: String?
37+
}
38+
39+
public init(accessKeyId: String, secretAccessKey: String,
40+
region: AWSRegion, endpointHostName: String,
41+
tableName: String,
42+
eventLoopProvider: HTTPClient.EventLoopProvider = .spawnNewThreads) {
43+
let staticCredentials = StaticCredentials(accessKeyId: accessKeyId,
44+
secretAccessKey: secretAccessKey,
45+
sessionToken: nil)
46+
47+
self.dynamodb = AWSDynamoDBClient(credentialsProvider: staticCredentials,
48+
awsRegion: region,
49+
endpointHostName: endpointHostName,
50+
eventLoopProvider: eventLoopProvider)
51+
self.targetTableName = tableName
52+
53+
Log.info("AWSDynamoDBTable created with region '\(region)' and hostname: '\(endpointHostName)'")
54+
}
55+
56+
public init(credentialsProvider: CredentialsProvider,
57+
region: AWSRegion, endpointHostName: String,
58+
tableName: String,
59+
eventLoopProvider: HTTPClient.EventLoopProvider = .spawnNewThreads) {
60+
self.dynamodb = AWSDynamoDBClient(credentialsProvider: credentialsProvider,
61+
awsRegion: region,
62+
endpointHostName: endpointHostName,
63+
eventLoopProvider: eventLoopProvider)
64+
self.targetTableName = tableName
65+
66+
Log.info("AWSDynamoDBTable created with region '\(region)' and hostname: '\(endpointHostName)'")
67+
}
68+
69+
/**
70+
Gracefully shuts down the client behind this table. This function is idempotent and
71+
will handle being called multiple times.
72+
*/
73+
public func close() {
74+
dynamodb.close()
75+
}
76+
77+
/**
78+
Waits for the client behind this table to be closed. If close() is not called,
79+
this will block forever.
80+
*/
81+
public func wait() {
82+
dynamodb.wait()
83+
}
84+
}

‎Sources/SmokeDynamoDB/AWSDynamoDBTable+DynamoDBTableAsync.swift

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -158,9 +158,10 @@ public extension AWSDynamoDBTable {
158158
limit: Int, exclusiveStartKey: String?,
159159
completion: @escaping (HTTPResult<([PolymorphicDatabaseItem<AttributesType, PossibleTypes>], String?)>) -> ())
160160
throws where AttributesType: PrimaryKeyAttributes, PossibleTypes: PossibleItemTypes {
161-
let queryInput = try getQueryInput(forPartitionKey: partitionKey, primaryKeyType: AttributesType.self,
162-
sortKeyCondition: sortKeyCondition, limit: limit,
163-
exclusiveStartKey: exclusiveStartKey)
161+
let queryInput = try DynamoDBModel.QueryInput.forSortKeyCondition(forPartitionKey: partitionKey, targetTableName: targetTableName,
162+
primaryKeyType: AttributesType.self,
163+
sortKeyCondition: sortKeyCondition, limit: limit,
164+
exclusiveStartKey: exclusiveStartKey)
164165
try dynamodb.queryAsync(input: queryInput) { result in
165166
switch result {
166167
case .response(let queryOutput):

‎Sources/SmokeDynamoDB/AWSDynamoDBTable+DynamoDBTableSync.swift

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -103,9 +103,10 @@ public extension AWSDynamoDBTable {
103103
exclusiveStartKey: String?) throws
104104
-> ([PolymorphicDatabaseItem<AttributesType, PossibleTypes>], String?)
105105
where AttributesType: PrimaryKeyAttributes, PossibleTypes: PossibleItemTypes {
106-
let queryInput = try getQueryInput(forPartitionKey: partitionKey, primaryKeyType: AttributesType.self,
107-
sortKeyCondition: sortKeyCondition, limit: limit,
108-
exclusiveStartKey: exclusiveStartKey)
106+
let queryInput = try DynamoDBModel.QueryInput.forSortKeyCondition(forPartitionKey: partitionKey, targetTableName: targetTableName,
107+
primaryKeyType: AttributesType.self,
108+
sortKeyCondition: sortKeyCondition, limit: limit,
109+
exclusiveStartKey: exclusiveStartKey)
109110
let queryOutput = try dynamodb.querySync(input: queryInput)
110111

111112
let lastEvaluatedKey: String?

0 commit comments

Comments
 (0)