PostgresDriver: 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

@@ -29,14 +29,14 @@ 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([$this, 'quoteIdentifier'], $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,8 +31,12 @@ class PostgresInsertQuery extends InsertQuery implements ReturningInterface
     /** @var PostgresDriver|null */
     protected ?DriverInterface $driver = null;
 
+    /** @deprecated */
     protected ?string $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(
@@ -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,
             '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;
 }