Postgres: add support for multiple returning columns

src/Driver/CompilerCache.php

@@ -89,7 +89,12 @@ public function compile(QueryParameters $params, string $prefix, FragmentInterfa
      */
     protected function hashInsertQuery(QueryParameters $params, array $tokens): string
     {
-        $hash = 'i_' . $tokens['table'] . implode('_', $tokens['columns']) . '_r' . ($tokens['return'] ?? '');
+        $hash = \sprintf(
+            'i_%s%s_r%s',
+            $tokens['table'],
+            \implode('_', $tokens['columns']),
+            \implode('_', (array)($tokens['return'] ?? []))
+        );
         foreach ($tokens['values'] as $value) {
             if ($value instanceof FragmentInterface) {
                 if ($value instanceof Expression || $value instanceof Fragment) {

src/Driver/Postgres/PostgresCompiler.php

@@ -14,6 +14,7 @@
 use Cycle\Database\Driver\CachingCompilerInterface;
 use Cycle\Database\Driver\Compiler;
 use Cycle\Database\Driver\Quoter;
+use Cycle\Database\Injection\FragmentInterface;
 use Cycle\Database\Injection\Parameter;
 use Cycle\Database\Query\QueryParameters;
 
@@ -29,14 +30,19 @@ protected function insertQuery(QueryParameters $params, Quoter $q, array $tokens
     {
         $result = parent::insertQuery($params, $q, $tokens);
 
-        if ($tokens['return'] === null) {
+        if (empty($tokens['return'])) {
             return $result;
         }
 
-        return sprintf(
+        return \sprintf(
             '%s RETURNING %s',
             $result,
-            $this->quoteIdentifier($tokens['return'])
+            \implode(',', \array_map(
+                fn (string|FragmentInterface|null $return) => $return instanceof FragmentInterface
+                    ? (string) $return
+                    : $this->quoteIdentifier($return),
+                $tokens['return']
+            ))
         );
     }
 

src/Driver/Postgres/PostgresDriver.php

@@ -98,7 +98,7 @@ public function shouldUseDefinedSchemas(): bool
     public function getPrimaryKey(string $prefix, string $table): ?string
     {
         $name = $prefix . $table;
-        if (array_key_exists($name, $this->primaryKeys)) {
+        if (\array_key_exists($name, $this->primaryKeys)) {
             return $this->primaryKeys[$name];
         }
 

src/Driver/Postgres/Query/PostgresInsertQuery.php

@@ -20,6 +20,7 @@
 use Cycle\Database\Query\InsertQuery;
 use Cycle\Database\Query\QueryInterface;
 use Cycle\Database\Query\QueryParameters;
+use Cycle\Database\StatementInterface;
 use Throwable;
 
 /**
@@ -30,7 +31,11 @@ class PostgresInsertQuery extends InsertQuery implements ReturningInterface
     /** @var PostgresDriver|null */
     protected ?DriverInterface $driver = null;
 
-    protected ?string $returning = 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
     {
@@ -48,13 +53,9 @@ public function returning(string|FragmentInterface ...$columns): self
     {
         $columns === [] and throw new BuilderException('RETURNING clause should contain at least 1 column.');
 
-        if (count($columns) > 1) {
-            throw new BuilderException(
-                'Postgres driver supports only single column returning at this moment.'
-            );
-        }
+        $this->returning = \count($columns) === 1 ? \reset($columns) : null;
 
-        $this->returning = (string)$columns[0];
+        $this->returningColumns = \array_values($columns);
 
         return $this;
     }
@@ -69,6 +70,15 @@ public function run(): mixed
         $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();
             }
@@ -83,23 +93,18 @@ public function getTokens(): array
     {
         return [
             'table' => $this->table,
-            'return' => $this->getPrimaryKey(),
+            'return' => $this->returningColumns !== [] ? $this->returningColumns : (array) $this->getPrimaryKey(),
             'columns' => $this->columns,
             'values' => $this->values,
         ];
     }
 
     private function getPrimaryKey(): ?string
     {
-        $primaryKey = $this->returning;
-        if ($primaryKey === null && $this->driver !== null && $this->table !== null) {
-            try {
-                $primaryKey = $this->driver->getPrimaryKey($this->prefix, $this->table);
-            } catch (Throwable) {
-                return null;
-            }
+        try {
+            return $this->driver?->getPrimaryKey($this->prefix, $this->table);
+        } catch (Throwable) {
+            return null;
         }
-
-        return $primaryKey;
     }
 }

src/Query/InsertQuery.php

@@ -112,8 +112,9 @@ public function values(mixed $rowsets): self
 
     /**
      * Run the query and return last insert id.
+     * Returns an assoc array of values if multiple columns were specified as returning columns.
      *
-     * @psalm-return int|non-empty-string|null
+     * @return array<non-empty-string, mixed>|int|non-empty-string|null
      */
     public function run(): mixed
     {

src/Query/ReturningInterface.php

@@ -4,12 +4,21 @@
 
 namespace Cycle\Database\Query;
 
+use Cycle\Database\Exception\BuilderException;
 use Cycle\Database\Injection\FragmentInterface;
 
 interface ReturningInterface extends QueryInterface
 {
     /**
      * Set returning column or expression.
+     *
+     * If set multiple columns and the driver supports it, then an insert result will be an array of values.
+     * If set one column and the driver supports it, then an insert result will be a single value,
+     * not an array of values.
+     *
+     * If set multiple columns and the driver does not support it, an exception will be thrown.
+     *
+     * @throws BuilderException
      */
     public function returning(string|FragmentInterface ...$columns): self;
 }

tests/Database/Functional/Driver/Postgres/Query/InsertQueryTest.php

@@ -6,6 +6,7 @@
 
 // phpcs:ignore
 use Cycle\Database\Driver\Postgres\Query\PostgresInsertQuery;
+use Cycle\Database\Driver\Postgres\Schema\PostgresColumn;
 use Cycle\Database\Exception\BuilderException;
 use Cycle\Database\Injection\Fragment;
 use Cycle\Database\Tests\Functional\Driver\Common\Query\InsertQueryTest as CommonClass;
@@ -91,39 +92,109 @@ public function testCustomReturning(): void
         );
     }
 
-    public function testCustomReturningWithFragment(): void
+    public function testCustomMultipleReturning(): void
     {
         $insert = $this->database->insert()->into('table')
             ->columns('name', 'balance')
             ->values('Anton', 100)
-            ->returning(new Fragment('COUNT(name)'));
+            ->returning('name', 'created_at');
 
         $this->assertSameQuery(
-            'INSERT INTO {table} ({name}, {balance}) VALUES (?, ?) RETURNING {COUNT(name)}',
+            'INSERT INTO {table} ({name}, {balance}) VALUES (?, ?) RETURNING {name}, {created_at}',
             $insert
         );
     }
 
-    public function testCustomReturningShouldContainColumns(): void
+    public function testCustomReturningWithFragment(): void
     {
-        $this->expectException(BuilderException::class);
-        $this->expectExceptionMessage('RETURNING clause should contain at least 1 column.');
+        $insert = $this->database->insert()->into('table')
+            ->columns('name', 'balance')
+            ->values('Anton', 100)
+            ->returning(new Fragment('"name" as "full_name"'));
 
-        $this->database->insert()->into('table')
+        $this->assertSameQuery(
+            'INSERT INTO {table} ({name}, {balance}) VALUES (?, ?) RETURNING {name} as {full_name}',
+            $insert
+        );
+    }
+
+    public function testCustomMultipleReturningWithFragment(): void
+    {
+        $insert = $this->database->insert()->into('table')
             ->columns('name', 'balance')
             ->values('Anton', 100)
-            ->returning();
+            ->returning('name', new Fragment('"created_at" as "date"'));
+
+        $this->assertSameQuery(
+            'INSERT INTO {table} ({name}, {balance}) VALUES (?, ?) RETURNING {name}, {created_at} as {date}',
+            $insert
+        );
+    }
+
+    public function testReturningValuesFromDatabase(): void
+    {
+        $schema = $this->schema('returning_values');
+        $schema->primary('id');
+        $schema->string('name');
+        $schema->serial('sort');
+        $schema->datetime('datetime', defaultValue: PostgresColumn::DATETIME_NOW);
+        $schema->save();
+
+        $returning = $this->database
+            ->insert('returning_values')
+            ->values(['name' => 'foo'])
+            ->returning('sort', 'datetime')
+            ->run();
+
+        $this->assertSame(1, $returning['sort']);
+        $this->assertIsString($returning['datetime']);
+        $this->assertNotFalse(\strtotime($returning['datetime']));
+
+        $returning = $this->database
+            ->insert('returning_values')
+            ->values(['name' => 'foo'])
+            ->returning('sort', new Fragment('"datetime" as "created_at"'))
+            ->run();
+
+        $this->assertSame(2, $returning['sort']);
+        $this->assertIsString($returning['created_at']);
+        $this->assertNotFalse(\strtotime($returning['created_at']));
     }
 
-    public function testCustomReturningSupportsOnlySingleColumn(): void
+    public function testReturningSingleValueFromDatabase(): void
+    {
+        $schema = $this->schema('returning_value');
+        $schema->primary('id');
+        $schema->string('name');
+        $schema->serial('sort');
+        $schema->save();
+
+        $returning = $this->database
+            ->insert('returning_value')
+            ->values(['name' => 'foo'])
+            ->returning('sort')
+            ->run();
+
+        $this->assertSame(1, $returning);
+
+        $returning = $this->database
+            ->insert('returning_value')
+            ->values(['name' => 'foo'])
+            ->returning(new Fragment('"sort" as "number"'))
+            ->run();
+
+        $this->assertSame(2, $returning);
+    }
+
+    public function testCustomReturningShouldContainColumns(): void
     {
         $this->expectException(BuilderException::class);
-        $this->expectExceptionMessage('Postgres driver supports only single column returning at this moment.');
+        $this->expectExceptionMessage('RETURNING clause should contain at least 1 column.');
 
         $this->database->insert()->into('table')
             ->columns('name', 'balance')
             ->values('Anton', 100)
-            ->returning('name', 'id');
+            ->returning();
     }
 
     public function testInsertMicroseconds(): void

tests/Database/Unit/Driver/CompilerCacheTest.php

@@ -0,0 +1,31 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Cycle\Database\Tests\Unit\Driver;
+
+use Cycle\Database\Driver\CompilerCache;
+use Cycle\Database\Driver\Postgres\PostgresCompiler;
+use Cycle\Database\Injection\Fragment;
+use Cycle\Database\Query\QueryParameters;
+use PHPUnit\Framework\TestCase;
+
+final class CompilerCacheTest extends TestCase
+{
+    public function testHashInsertQueryWithReturningFragment(): void
+    {
+        $compiler = new CompilerCache(new PostgresCompiler());
+        $ref = new \ReflectionMethod($compiler, 'hashInsertQuery');
+        $ref->setAccessible(true);
+
+        $this->assertSame(
+            'i_some_tablename_full_name_rname_"full_name" as "fullName"P?',
+            $ref->invoke($compiler, new QueryParameters(), [
+                'table' => 'some_table',
+                'columns' => ['name', 'full_name'],
+                'values' => ['Foo'],
+                'return' => ['name', new Fragment('"full_name" as "fullName"')],
+            ])
+        );
+    }
+}