Merge pull request #117 from lifion/bugfix/add-shards-to-enhanced-con…

…sumers Add shards on update
lifion · Apr 29, 2019 · fab8c17 · fab8c17
2 parents f814f18 + 9d4a7d7
commit fab8c17
Showing 1 changed file with 29 additions and 13 deletions.
diff --git a/lib/state-store.js b/lib/state-store.js
@@ -459,23 +459,39 @@ class StateStore {
   }
 
   async lockStreamConsumer(consumerName, version) {
-    const { client, consumerGroup, consumerId, logger, streamName } = internal(this);
+    const {
+      client,
+      consumerGroup,
+      consumerId,
+      logger,
+      streamName,
+      useAutoShardAssignment
+    } = internal(this);
+
     try {
       await client.update({
         ConditionExpression: `#a.#b.#d = :z`,
-        ExpressionAttributeNames: {
-          '#a': 'enhancedConsumers',
-          '#b': consumerName,
-          '#c': 'isUsedBy',
-          '#d': 'version'
-        },
-        ExpressionAttributeValues: {
-          ':x': consumerId,
-          ':y': generate(),
-          ':z': version
-        },
+        ExpressionAttributeNames: Object.assign(
+          {
+            '#a': 'enhancedConsumers',
+            '#b': consumerName,
+            '#c': 'isUsedBy',
+            '#d': 'version'
+          },
+          !useAutoShardAssignment && { '#e': 'shards' }
+        ),
+        ExpressionAttributeValues: Object.assign(
+          {
+            ':x': consumerId,
+            ':y': generate(),
+            ':z': version
+          },
+          !useAutoShardAssignment && { ':e': {} }
+        ),
         Key: { consumerGroup, streamName },
-        UpdateExpression: `SET #a.#b.#c = :x, #a.#b.#d = :y`
+        UpdateExpression: `SET #a.#b.#c = :x, #a.#b.#d = :y${
+          !useAutoShardAssignment ? ', #a.#b.#e = if_not_exists(#a.#b.#e, :e)' : ''
+        }`
       });
       return true;
     } catch (err) {