From 3e6dee36e112ab9592a98334be412071d605623c Mon Sep 17 00:00:00 2001
From: Alexander Dylevskiy <alexd@route4me.com>
Date: Thu, 1 Jun 2023 10:15:34 +0300
Subject: [PATCH] feat: Add insertOrUpdateUsingMutation

+ change ManagesMutations private methods to protected.
---
 CHANGELOG.md                         |  1 +
 README.md                            |  1 +
 src/Concerns/ManagesMutations.php    | 20 +++++++++++--
 src/Query/Concerns/UsesMutations.php |  9 ++++++
 tests/ConnectionTest.php             | 45 ++++++++++++++++++++++++++++
 5 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8ec54098..760e98d9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 Added
 - Added deprecation warnings to `Connection::runDdl` and `Connection::runDdls` (#98)
+- Added  `ManagesMutations::insertOrUpdateUsingMutation` and `UsesMutations::insertOrUpdateUsingMutation` to do upserts (#109)
 
 Changed
 - `Connection::waitForOperation` and `Connection::isDoneOperation` has been removed. (#99)
diff --git a/README.md b/README.md
index a52957de..42df3a53 100644
--- a/README.md
+++ b/README.md
@@ -258,6 +258,7 @@ You can [insert, update, and delete data using mutations](https://cloud.google.c
 ```
 $queryBuilder->insertUsingMutation($values);
 $queryBuilder->updateUsingMutation($values);
+$queryBuilder->insertOrUpdateUsingMutation($values);
 $queryBuilder->deleteUsingMutation($values);
 ```
 
diff --git a/src/Concerns/ManagesMutations.php b/src/Concerns/ManagesMutations.php
index 5a25c2bb..5e96c714 100644
--- a/src/Concerns/ManagesMutations.php
+++ b/src/Concerns/ManagesMutations.php
@@ -59,6 +59,20 @@ public function updateUsingMutation(string $table, array $dataSet)
         });
     }
 
+    /**
+     * @param string $table
+     * @param array<string, mixed> $dataSet
+     * @return void
+     */
+    public function insertOrUpdateUsingMutation(string $table, array $dataSet)
+    {
+        $this->withTransactionEvents(function () use ($table, $dataSet) {
+            $dataSet = $this->prepareForMutation($dataSet);
+            $this->event(new MutatingData($this, $table, 'update', $dataSet));
+            $this->getDatabaseContext()->insertOrUpdateBatch($table, $dataSet);
+        });
+    }
+
     /**
      * @param string $table
      * @param scalar|array<mixed>|KeySet $keySet
@@ -78,7 +92,7 @@ public function deleteUsingMutation(string $table, $keySet)
      * @param callable $mutationCall
      * @return void
      */
-    private function withTransactionEvents(callable $mutationCall)
+    protected function withTransactionEvents(callable $mutationCall)
     {
         // events not necessary since it is already called
         if ($this->inTransaction()) {
@@ -94,7 +108,7 @@ private function withTransactionEvents(callable $mutationCall)
      * @param array $dataSet
      * @return array
      */
-    private function prepareForMutation(array $dataSet): array
+    protected function prepareForMutation(array $dataSet): array
     {
         if (empty($dataSet)) {
             return [];
@@ -119,7 +133,7 @@ private function prepareForMutation(array $dataSet): array
      * @param mixed|list<string>|KeySet $keys
      * @return KeySet
      */
-    private function createDeleteMutationKeySet($keys)
+    protected function createDeleteMutationKeySet($keys)
     {
         if ($keys instanceof KeySet) {
             return $keys;
diff --git a/src/Query/Concerns/UsesMutations.php b/src/Query/Concerns/UsesMutations.php
index b50f0a0d..7d1afbf4 100644
--- a/src/Query/Concerns/UsesMutations.php
+++ b/src/Query/Concerns/UsesMutations.php
@@ -43,6 +43,15 @@ public function updateUsingMutation(array $values)
         $this->connection->updateUsingMutation($this->from, $values);
     }
 
+    /**
+     * @param array<string, mixed> $values
+     * @return void
+     */
+    public function insertOrUpdateUsingMutation(array $values)
+    {
+        $this->connection->insertOrUpdateUsingMutation($this->from, $values);
+    }
+
     /**
      * @param array|KeySet $keys
      * @return void
diff --git a/tests/ConnectionTest.php b/tests/ConnectionTest.php
index 2d656b7a..928c420a 100644
--- a/tests/ConnectionTest.php
+++ b/tests/ConnectionTest.php
@@ -137,6 +137,51 @@ public function testUpdateUsingMutationWithoutTransaction(): void
         Event::assertDispatchedTimes(TransactionCommitted::class, 2);
     }
 
+    public function testInsertOrUpdateUsingMutationWithTransaction(): void
+    {
+        Event::fake();
+
+        $userId1 = $this->generateUuid();
+        $userId2 = $this->generateUuid();
+        $conn = $this->getDefaultConnection();
+
+        $conn->insertUsingMutation(self::TABLE_NAME_USER, ['userId' => $userId1, 'name' => 'test1']);
+
+        $conn->transaction(function () use ($conn, $userId1, $userId2) {
+            $conn->insertOrUpdateUsingMutation(self::TABLE_NAME_USER, [
+                ['userId' => $userId1, 'name' => 'tester1'],
+                ['userId' => $userId2, 'name' => 'tester2'],
+            ]);
+        });
+
+        $this->assertEquals(['userId' => $userId1, 'name' => 'tester1'], $conn->table(self::TABLE_NAME_USER)->where('userId', $userId1)->first());
+        $this->assertEquals(['userId' => $userId2, 'name' => 'tester2'], $conn->table(self::TABLE_NAME_USER)->where('userId', $userId2)->first());
+        Event::assertDispatchedTimes(TransactionBeginning::class, 2);
+        Event::assertDispatchedTimes(MutatingData::class, 2);
+        Event::assertDispatchedTimes(TransactionCommitted::class, 2);
+    }
+
+    public function testInsertOrUpdateUsingMutationWithoutTransaction(): void
+    {
+        Event::fake();
+
+        $userId1 = $this->generateUuid();
+        $userId2 = $this->generateUuid();
+        $conn = $this->getDefaultConnection();
+
+        $conn->insertUsingMutation(self::TABLE_NAME_USER, ['userId' => $userId1, 'name' => 'test1']);
+        $conn->insertOrUpdateUsingMutation(self::TABLE_NAME_USER, [
+            ['userId' => $userId1, 'name' => 'tester1'],
+            ['userId' => $userId2, 'name' => 'tester2'],
+        ]);
+
+        $this->assertEquals(['userId' => $userId1, 'name' => 'tester1'], $conn->table(self::TABLE_NAME_USER)->where('userId', $userId1)->first());
+        $this->assertEquals(['userId' => $userId2, 'name' => 'tester2'], $conn->table(self::TABLE_NAME_USER)->where('userId', $userId2)->first());
+        Event::assertDispatchedTimes(TransactionBeginning::class, 2);
+        Event::assertDispatchedTimes(MutatingData::class, 2);
+        Event::assertDispatchedTimes(TransactionCommitted::class, 2);
+    }
+
     public function testDeleteUsingMutationWithTransaction(): void
     {
         Event::fake();