Skip to content

Commit f402aa5

Browse files
committed
Support for rawsource seeds that use the AWS AttributeValues instead of Javascript types.
1 parent 168400d commit f402aa5

File tree

2 files changed

+42
-12
lines changed

2 files changed

+42
-12
lines changed

index.js

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -153,15 +153,25 @@ class ServerlessDynamodbLocal {
153153

154154
seedHandler() {
155155
const options = this.options;
156-
const documentClient = this.dynamodbOptions(options).doc;
157-
const seedSources = this.seedSources;
158-
return BbPromise.each(seedSources, (source) => {
156+
const dynamodb = this.dynamodbOptions(options);
157+
158+
const seedPromise = BbPromise.each(this.seedSources, (source) => {
159159
if (!source.table) {
160160
throw new Error("seeding source \"table\" property not defined");
161161
}
162162
return seeder.locateSeeds(source.sources || [])
163-
.then((seeds) => seeder.writeSeeds(documentClient, source.table, seeds));
163+
.then((seeds) => seeder.writeSeeds(dynamodb.doc.batchWrite.bind(dynamodb.doc), source.table, seeds));
164+
});
165+
166+
const rawSeedPromise = BbPromise.each(this.rawSeedSources, (source) => {
167+
if (!source.table) {
168+
throw new Error("Raw seeding source \"table\" property not defined");
169+
}
170+
return seeder.locateSeeds(source.rawsources || [])
171+
.then((seeds) => seeder.writeSeeds(dynamodb.raw.batchWriteItem.bind(dynamodb.raw), source.table, seeds));
164172
});
173+
174+
return BbPromise.all([seedPromise, rawSeedPromise]);
165175
}
166176

167177
removeHandler() {
@@ -258,6 +268,26 @@ class ServerlessDynamodbLocal {
258268
return [].concat.apply([], sourcesByCategory);
259269
}
260270

271+
/**
272+
* Gets the raw seeding sources
273+
*/
274+
get rawSeedSources() {
275+
const config = this.service.custom.dynamodb;
276+
const seedConfig = _.get(config, "seed", {});
277+
const seed = this.options.seed || config.start.seed || seedConfig;
278+
let categories;
279+
if (typeof seed === "string") {
280+
categories = seed.split(",");
281+
} else if(seed) {
282+
categories = Object.keys(seedConfig);
283+
} else { // if (!seed)
284+
this.serverlessLog("DynamoDB - No seeding defined. Skipping data seeding.");
285+
return [];
286+
}
287+
const sourcesByCategory = categories.map((category) => seedConfig[category].rawsources);
288+
return [].concat.apply([], sourcesByCategory);
289+
}
290+
261291
createTable(dynamodb, migration) {
262292
return new BbPromise((resolve, reject) => {
263293
if (migration.StreamSpecification && migration.StreamSpecification.StreamViewType) {

src/seeder.js

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,11 @@ const MIGRATION_SEED_CONCURRENCY = 5;
1515
/**
1616
* Writes a batch chunk of migration seeds to DynamoDB. DynamoDB has a limit on the number of
1717
* items that may be written in a batch operation.
18-
* @param {DynamoDocumentClient} dynamodb The DynamoDB Document client
18+
* @param {function} dynamodbWriteFunction The DynamoDB DocumentClient.batchWrite or DynamoDB.batchWriteItem function
1919
* @param {string} tableName The table name being written to
2020
* @param {any[]} seeds The migration seeds being written to the table
2121
*/
22-
function writeSeedBatch(dynamodb, tableName, seeds) {
22+
function writeSeedBatch(dynamodbWriteFunction, tableName, seeds) {
2323
const params = {
2424
RequestItems: {
2525
[tableName]: seeds.map((seed) => ({
@@ -34,7 +34,7 @@ function writeSeedBatch(dynamodb, tableName, seeds) {
3434
// again a few times in case the Database resources are in the middle of provisioning.
3535
let interval = 0;
3636
function execute(interval) {
37-
setTimeout(() => dynamodb.batchWrite(params, (err) => {
37+
setTimeout(() => dynamodbWriteFunction(params, (err) => {
3838
if (err) {
3939
if (err.code === "ResourceNotFoundException" && interval <= 5000) {
4040
execute(interval + 1000);
@@ -52,13 +52,13 @@ function writeSeedBatch(dynamodb, tableName, seeds) {
5252

5353
/**
5454
* Writes a seed corpus to the given database table
55-
* @param {DocumentClient} dynamodb The DynamoDB document instance
55+
* @param {function} dynamodbWriteFunction The DynamoDB DocumentClient.batchWrite or DynamoDB.batchWriteItem function
5656
* @param {string} tableName The table name
5757
* @param {any[]} seeds The seed values
5858
*/
59-
function writeSeeds(dynamodb, tableName, seeds) {
60-
if (!dynamodb) {
61-
throw new Error("dynamodb argument must be provided");
59+
function writeSeeds(dynamodbWriteFunction, tableName, seeds) {
60+
if (!dynamodbWriteFunction) {
61+
throw new Error("dynamodbWriteFunction argument must be provided");
6262
}
6363
if (!tableName) {
6464
throw new Error("table name argument must be provided");
@@ -71,7 +71,7 @@ function writeSeeds(dynamodb, tableName, seeds) {
7171
const seedChunks = _.chunk(seeds, MAX_MIGRATION_CHUNK);
7272
return BbPromise.map(
7373
seedChunks,
74-
(chunk) => writeSeedBatch(dynamodb, tableName, chunk),
74+
(chunk) => writeSeedBatch(dynamodbWriteFunction, tableName, chunk),
7575
{ concurrency: MIGRATION_SEED_CONCURRENCY }
7676
)
7777
.then(() => console.log("Seed running complete for table: " + tableName));

0 commit comments

Comments
 (0)