Upsert feature

.gitignore

@@ -1,4 +1,5 @@
 /.idea/
+/.report/
 /.dolt/
 /runtime/
 /vendor/

composer.json

@@ -93,6 +93,8 @@
         "cs:fix": "php-cs-fixer fix -v",
         "psalm": "psalm",
         "psalm:baseline": "psalm --set-baseline=psalm-baseline.xml",
-        "test": "phpunit --color=always"
+        "test": "phpunit --color=always",
+        "test:coverage": "phpunit --colors=always --coverage-text",
+        "test:report": "phpunit --do-not-cache-result --log-junit=.report/phpunit.xml --coverage-clover=.report/coverage-clover.xml --coverage-cobertura=.report/coverage-cobertura.xml --coverage-text --colors=never"
     }
 }

src/Database.php

@@ -17,6 +17,7 @@
 use Cycle\Database\Query\InsertQuery;
 use Cycle\Database\Query\SelectQuery;
 use Cycle\Database\Query\UpdateQuery;
+use Cycle\Database\Query\UpsertQuery;
 
 /**
  * Database class is high level abstraction at top of Driver. Databases usually linked to real
@@ -139,6 +140,13 @@ public function insert(?string $table = null): InsertQuery
             ->insertQuery($this->prefix, $table);
     }
 
+    public function upsert(?string $table = null): UpsertQuery
+    {
+        return $this->getDriver(self::WRITE)
+            ->getQueryBuilder()
+            ->upsertQuery($this->prefix, $table);
+    }
+
     public function update(?string $table = null, array $values = [], array $where = []): UpdateQuery
     {
         return $this->getDriver(self::WRITE)

src/DatabaseInterface.php

@@ -17,6 +17,7 @@
 use Cycle\Database\Query\InsertQuery;
 use Cycle\Database\Query\SelectQuery;
 use Cycle\Database\Query\UpdateQuery;
+use Cycle\Database\Query\UpsertQuery;
 
 /**
  * DatabaseInterface is high level abstraction used to represent single database. You must always
@@ -104,6 +105,15 @@ public function query(string $query, array $parameters = []): StatementInterface
      */
     public function insert(string $table = ''): InsertQuery;
 
+    /**
+     * Get instance of UpsertBuilder associated with current Database.
+     *
+     * @param string $table Table where values should be upserted to.
+     *
+     * @see self::withoutCache() May be useful to disable query cache for batch inserts.
+     */
+    public function upsert(string $table = ''): UpsertQuery;
+
     /**
      * Get instance of UpdateBuilder associated with current Database.
      *

src/Driver/Compiler.php

@@ -109,6 +109,9 @@ protected function fragment(
             case self::INSERT_QUERY:
                 return $this->insertQuery($params, $q, $tokens);
 
+            case self::UPSERT_QUERY:
+                return $this->upsertQuery($params, $q, $tokens);
+
             case self::SELECT_QUERY:
                 if ($nestedQuery) {
                     if ($fragment->getPrefix() !== null) {
@@ -169,6 +172,43 @@ protected function insertQuery(QueryParameters $params, Quoter $q, array $tokens
         );
     }
 
+    /**
+     * @psalm-return non-empty-string
+     */
+    protected function upsertQuery(QueryParameters $params, Quoter $q, array $tokens): string
+    {
+        if (\count($tokens['conflicts']) === 0) {
+            throw new CompilerException('Upsert query must define conflicting index column names');
+        }
+
+        if (\count($tokens['columns']) === 0) {
+            throw new CompilerException('Upsert query must define at least one column');
+        }
+
+        $values = [];
+
+        foreach ($tokens['values'] as $value) {
+            $values[] = $this->value($params, $q, $value);
+        }
+
+        $updates = \array_map(
+            function (string $column) use ($params, $q) {
+                $name   = $this->name($params, $q, $column);
+                return \sprintf('%s = EXCLUDED.%s', $name, $name);
+            },
+            $tokens['columns'],
+        );
+
+        return \sprintf(
+            'INSERT INTO %s (%s) VALUES %s ON CONFLICT (%s) DO UPDATE SET %s',
+            $this->name($params, $q, $tokens['table'], true),
+            $this->columns($params, $q, $tokens['columns']),
+            \implode(', ', $values),
+            $this->columns($params, $q, $tokens['conflicts']),
+            \implode(', ', $updates),
+        );
+    }
+
     /**
      * @psalm-return non-empty-string
      */

src/Driver/CompilerInterface.php

@@ -25,6 +25,7 @@ interface CompilerInterface
     public const DELETE_QUERY = 7;
     public const JSON_EXPRESSION = 8;
     public const SUBQUERY = 9;
+    public const UPSERT_QUERY = 10;
     public const TOKEN_AND = '@AND';
     public const TOKEN_OR = '@OR';
     public const TOKEN_AND_NOT = '@AND NOT';

src/Driver/MySQL/MySQLCompiler.php

@@ -15,6 +15,7 @@
 use Cycle\Database\Driver\Compiler;
 use Cycle\Database\Driver\MySQL\Injection\CompileJson;
 use Cycle\Database\Driver\Quoter;
+use Cycle\Database\Exception\CompilerException;
 use Cycle\Database\Injection\FragmentInterface;
 use Cycle\Database\Injection\Parameter;
 use Cycle\Database\Query\QueryParameters;
@@ -36,6 +37,38 @@ protected function insertQuery(QueryParameters $params, Quoter $q, array $tokens
         return parent::insertQuery($params, $q, $tokens);
     }
 
+    /**
+     * @psalm-return non-empty-string
+     */
+    protected function upsertQuery(QueryParameters $params, Quoter $q, array $tokens): string
+    {
+        if (\count($tokens['columns']) === 0) {
+            throw new CompilerException('Upsert query must define at least one column');
+        }
+
+        $values = [];
+
+        foreach ($tokens['values'] as $value) {
+            $values[] = $this->value($params, $q, $value);
+        }
+
+        $updates = \array_map(
+            function ($column) use ($params, $q) {
+                $name   = $this->name($params, $q, $column);
+                return \sprintf('%s = VALUES(%s)', $name, $name);
+            },
+            $tokens['columns'],
+        );
+
+        return \sprintf(
+            'INSERT INTO %s (%s) VALUES %s ON DUPLICATE KEY UPDATE %s',
+            $this->name($params, $q, $tokens['table'], true),
+            $this->columns($params, $q, $tokens['columns']),
+            \implode(', ', $values),
+            \implode(', ', $updates),
+        );
+    }
+
     /**
      *
      *

src/Driver/MySQL/MySQLDriver.php

@@ -20,6 +20,7 @@
 use Cycle\Database\Exception\StatementException;
 use Cycle\Database\Query\InsertQuery;
 use Cycle\Database\Query\QueryBuilder;
+use Cycle\Database\Query\UpsertQuery;
 
 /**
  * Talks to mysql databases.
@@ -38,6 +39,7 @@ public static function create(DriverConfig $config): static
             new QueryBuilder(
                 new MySQLSelectQuery(),
                 new InsertQuery(),
+                new UpsertQuery(),
                 new MySQLUpdateQuery(),
                 new MySQLDeleteQuery(),
             ),

src/Driver/Postgres/PostgresCompiler.php

@@ -52,6 +52,29 @@ protected function insertQuery(QueryParameters $params, Quoter $q, array $tokens
         );
     }
 
+    /**
+     * @psalm-return non-empty-string
+     */
+    protected function upsertQuery(QueryParameters $params, Quoter $q, array $tokens): string
+    {
+        $query = parent::upsertQuery($params, $q, $tokens);
+
+        if (empty($tokens['return'])) {
+            return $query;
+        }
+
+        return \sprintf(
+            '%s RETURNING %s',
+            $query,
+            \implode(',', \array_map(
+                fn(string|FragmentInterface|null $return) => $return instanceof FragmentInterface
+                    ? $this->fragment($params, $q, $return)
+                    : $this->quoteIdentifier($return),
+                $tokens['return'],
+            )),
+        );
+    }
+
     protected function distinct(QueryParameters $params, Quoter $q, string|bool|array $distinct): string
     {
         if ($distinct === false) {

src/Driver/Postgres/PostgresDriver.php

@@ -19,6 +19,7 @@
 use Cycle\Database\Driver\Postgres\Query\PostgresInsertQuery;
 use Cycle\Database\Driver\Postgres\Query\PostgresSelectQuery;
 use Cycle\Database\Driver\Postgres\Query\PostgresUpdateQuery;
+use Cycle\Database\Driver\Postgres\Query\PostgresUpsertQuery;
 use Cycle\Database\Exception\DriverException;
 use Cycle\Database\Exception\StatementException;
 use Cycle\Database\Query\QueryBuilder;
@@ -65,6 +66,7 @@ public static function create(DriverConfig $config): static
             new QueryBuilder(
                 new PostgresSelectQuery(),
                 new PostgresInsertQuery(),
+                new PostgresUpsertQuery(),
                 new PostgresUpdateQuery(),
                 new PostgresDeleteQuery(),
             ),

src/Driver/Postgres/Query/PostgresUpsertQuery.php

@@ -0,0 +1,106 @@
+<?php
+
+/**
+ * This file is part of Cycle ORM package.
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+declare(strict_types=1);
+
+namespace Cycle\Database\Driver\Postgres\Query;
+
+use Cycle\Database\Driver\DriverInterface;
+use Cycle\Database\Driver\Postgres\PostgresDriver;
+use Cycle\Database\Exception\BuilderException;
+use Cycle\Database\Exception\ReadonlyConnectionException;
+use Cycle\Database\Injection\FragmentInterface;
+use Cycle\Database\Query\ReturningInterface;
+use Cycle\Database\Query\QueryInterface;
+use Cycle\Database\Query\QueryParameters;
+use Cycle\Database\Query\UpsertQuery;
+use Cycle\Database\StatementInterface;
+
+/**
+ * Postgres driver requires a slightly different way to handle last insert id.
+ */
+class PostgresUpsertQuery extends UpsertQuery implements ReturningInterface
+{
+    /** @var PostgresDriver|null */
+    protected ?DriverInterface $driver = null;
+
+    /** @deprecated */
+    protected string|FragmentInterface|null $returning = null;
+
+    /** @var list<FragmentInterface|non-empty-string> */
+    protected array $returningColumns = [];
+
+    public function withDriver(DriverInterface $driver, ?string $prefix = null): QueryInterface
+    {
+        $driver instanceof PostgresDriver or throw new BuilderException(
+            'Postgres UpsertQuery can be used only with Postgres driver',
+        );
+
+        return parent::withDriver($driver, $prefix);
+    }
+
+    /**
+     * Set returning column. If not set, the driver will detect PK automatically.
+     */
+    public function returning(string|FragmentInterface ...$columns): self
+    {
+        $columns === [] and throw new BuilderException('RETURNING clause should contain at least 1 column.');
+
+        $this->returning = \count($columns) === 1 ? \reset($columns) : null;
+
+        $this->returningColumns = \array_values($columns);
+
+        return $this;
+    }
+
+    public function run(): mixed
+    {
+        $params = new QueryParameters();
+        $queryString = $this->sqlStatement($params);
+
+        $this->driver->isReadonly() and throw ReadonlyConnectionException::onWriteStatementExecution();
+
+        $result = $this->driver->query($queryString, $params->getParameters());
+
+        try {
+            if ($this->returningColumns !== []) {
+                if (\count($this->returningColumns) === 1) {
+                    return $result->fetchColumn();
+                }
+
+                return $result->fetch(StatementInterface::FETCH_ASSOC);
+            }
+
+            // Return PK if no RETURNING clause is set
+            if ($this->getPrimaryKey() !== null) {
+                return $result->fetchColumn();
+            }
+
+            return null;
+        } finally {
+            $result->close();
+        }
+    }
+
+    public function getTokens(): array
+    {
+        return parent::getTokens() + [
+            'return'  => $this->returningColumns !== [] ? $this->returningColumns : (array) $this->getPrimaryKey(),
+        ];
+    }
+
+    private function getPrimaryKey(): ?string
+    {
+        try {
+            return $this->driver?->getPrimaryKey($this->prefix, $this->table);
+        } catch (\Throwable) {
+            return null;
+        }
+    }
+}