From 04d20781258a40953df0f6646c5624b7affb95b9 Mon Sep 17 00:00:00 2001
From: Bao Pham <gbaopham@gmail.com>
Date: Mon, 25 Jun 2018 20:17:52 +0700
Subject: [PATCH] Fix toDynamoDbQuery() together with decorate()

---
 src/DynamoDbModel.php        |  2 +-
 src/DynamoDbQueryBuilder.php |  5 +++++
 tests/DynamoDbModelTest.php  | 31 +++++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/DynamoDbModel.php b/src/DynamoDbModel.php
index 3cc5afd..e908c89 100644
--- a/src/DynamoDbModel.php
+++ b/src/DynamoDbModel.php
@@ -304,7 +304,7 @@ public function getKeys()
     }
 
     /**
-     * Get the primary for the model.
+     * Get the primary key for the model.
      *
      * @return string
      */
diff --git a/src/DynamoDbQueryBuilder.php b/src/DynamoDbQueryBuilder.php
index 512fae3..f4a8a66 100644
--- a/src/DynamoDbQueryBuilder.php
+++ b/src/DynamoDbQueryBuilder.php
@@ -674,6 +674,11 @@ public function toDynamoDbQuery(
         }
 
         $raw->query = $this->cleanUpQuery($query);
+
+        if ($this->decorator) {
+            call_user_func($this->decorator, $raw);
+        }
+
         return $raw;
     }
 
diff --git a/tests/DynamoDbModelTest.php b/tests/DynamoDbModelTest.php
index 7f34bb5..10541e9 100644
--- a/tests/DynamoDbModelTest.php
+++ b/tests/DynamoDbModelTest.php
@@ -779,6 +779,37 @@ public function testDecorateRawQuery()
         $this->assertEquals(range(1, 9), $items->pluck('count')->sort()->values()->toArray());
     }
 
+    public function testToDynamoDbQueryWithDecorate()
+    {
+        $queryWithoutDecorate = $this->testModel->where('foo', 'bar');
+
+        $this->assertEquals(
+            [
+                'FilterExpression' => '#foo = :a1',
+                'ExpressionAttributeNames' => ['#foo' => 'foo'],
+                'ExpressionAttributeValues' => [':a1' => ['S' => 'bar']],
+                'TableName' => $this->testModel->getTable(),
+            ],
+            $queryWithoutDecorate->toDynamoDbQuery()->query
+        );
+
+        $queryWithDecorate = $queryWithoutDecorate->clone()
+            ->decorate(function (RawDynamoDbQuery $raw) {
+                $raw->query['FilterExpression'] .= ' AND extra_col = :extra_col_val';
+                $raw->query['ExpressionAttributeValues'][':extra_col_val'] = ['N' => 0];
+            });
+
+        $this->assertEquals(
+            [
+                'FilterExpression' => '#foo = :a1 AND extra_col = :extra_col_val',
+                'ExpressionAttributeNames' => ['#foo' => 'foo'],
+                'ExpressionAttributeValues' => [':a1' => ['S' => 'bar'], ':extra_col_val' => ['N' => 0]],
+                'TableName' => $this->testModel->getTable(),
+            ],
+            $queryWithDecorate->toDynamoDbQuery()->query
+        );
+    }
+
     private function assertUsingKeyAndFilterConditions($model)
     {
         foreach (range(0, 9) as $i) {