Support multiple returning columns in MSSQL driver

src/Driver/SQLServer/Query/SQLServerInsertQuery.php

@@ -0,0 +1,84 @@
+<?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\SQLServer\Query;
+
+use Cycle\Database\Driver\DriverInterface;
+use Cycle\Database\Driver\SQLServer\SQLServerDriver;
+use Cycle\Database\Exception\BuilderException;
+use Cycle\Database\Exception\ReadonlyConnectionException;
+use Cycle\Database\Injection\FragmentInterface;
+use Cycle\Database\Query\QueryParameters;
+use Cycle\Database\Query\ReturningInterface;
+use Cycle\Database\Query\InsertQuery;
+use Cycle\Database\Query\QueryInterface;
+use Cycle\Database\StatementInterface;
+
+class SQLServerInsertQuery extends InsertQuery implements ReturningInterface
+{
+    /**
+     * @var SQLServerDriver|null
+     */
+    protected ?DriverInterface $driver = null;
+
+    /**
+     * @var list<FragmentInterface|non-empty-string>
+     */
+    protected array $returningColumns = [];
+
+    public function withDriver(DriverInterface $driver, string $prefix = null): QueryInterface
+    {
+        $driver instanceof SQLServerDriver or throw new BuilderException(
+            'SQLServer InsertQuery can be used only with SQLServer driver'
+        );
+
+        return parent::withDriver($driver, $prefix);
+    }
+
+    public function returning(string|FragmentInterface ...$columns): self
+    {
+        $columns === [] and throw new BuilderException('RETURNING clause should contain at least 1 column.');
+
+        $this->returningColumns = \array_values($columns);
+
+        return $this;
+    }
+
+    public function run(): mixed
+    {
+        if ($this->returningColumns === []) {
+            return parent::run();
+        }
+
+        $params = new QueryParameters();
+        $queryString = $this->sqlStatement($params);
+
+        $this->driver->isReadonly() and throw ReadonlyConnectionException::onWriteStatementExecution();
+
+        $result = $this->driver->query($queryString, $params->getParameters());
+
+        try {
+            if (\count($this->returningColumns) === 1) {
+                return $result->fetchColumn();
+            }
+            return $result->fetch(StatementInterface::FETCH_ASSOC);
+        } finally {
+            $result->close();
+        }
+    }
+
+    public function getTokens(): array
+    {
+        return parent::getTokens() + [
+            'return' => $this->returningColumns,
+        ];
+    }
+}

src/Driver/SQLServer/SQLServerCompiler.php

@@ -14,6 +14,7 @@
 use Cycle\Database\Driver\Compiler;
 use Cycle\Database\Driver\Quoter;
 use Cycle\Database\Injection\Fragment;
+use Cycle\Database\Injection\FragmentInterface;
 use Cycle\Database\Injection\Parameter;
 use Cycle\Database\Query\QueryParameters;
 
@@ -28,6 +29,44 @@ class SQLServerCompiler extends Compiler
      */
     public const ROW_NUMBER = '_ROW_NUMBER_';
 
+    /**
+     * @psalm-return non-empty-string
+     */
+    protected function insertQuery(QueryParameters $params, Quoter $q, array $tokens): string
+    {
+        if (empty($tokens['return'])) {
+            return parent::insertQuery($params, $q, $tokens);
+        }
+
+        $values = [];
+        foreach ($tokens['values'] as $value) {
+            $values[] = $this->value($params, $q, $value);
+        }
+
+        $output = \implode(',', \array_map(
+            fn (string|FragmentInterface|null $return) => $return instanceof FragmentInterface
+                ? (string) $return
+                : 'INSERTED.' . $this->quoteIdentifier($return),
+            $tokens['return']
+        ));
+
+        if ($tokens['columns'] === []) {
+            return \sprintf(
+                'INSERT INTO %s OUTPUT %s DEFAULT VALUES',
+                $this->name($params, $q, $tokens['table'], true),
+                $output
+            );
+        }
+
+        return \sprintf(
+            'INSERT INTO %s (%s) OUTPUT %s VALUES %s',
+            $this->name($params, $q, $tokens['table'], true),
+            $this->columns($params, $q, $tokens['columns']),
+            $output,
+            \implode(', ', $values)
+        );
+    }
+
     /**
      * {@inheritdoc}
      *

src/Driver/SQLServer/SQLServerDriver.php

@@ -17,12 +17,12 @@
 use Cycle\Database\Driver\Driver;
 use Cycle\Database\Driver\PDOStatementInterface;
 use Cycle\Database\Driver\SQLServer\Query\SQLServerDeleteQuery;
+use Cycle\Database\Driver\SQLServer\Query\SQLServerInsertQuery;
 use Cycle\Database\Driver\SQLServer\Query\SQLServerSelectQuery;
 use Cycle\Database\Driver\SQLServer\Query\SQLServerUpdateQuery;
 use Cycle\Database\Exception\DriverException;
 use Cycle\Database\Exception\StatementException;
 use Cycle\Database\Injection\ParameterInterface;
-use Cycle\Database\Query\InsertQuery;
 use Cycle\Database\Query\QueryBuilder;
 use PDO;
 
@@ -175,7 +175,7 @@ public static function create(DriverConfig $config): static
             new SQLServerCompiler('[]'),
             new QueryBuilder(
                 new SQLServerSelectQuery(),
-                new InsertQuery(),
+                new SQLServerInsertQuery(),
                 new SQLServerUpdateQuery(),
                 new SQLServerDeleteQuery()
             )

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

@@ -5,13 +5,164 @@
 namespace Cycle\Database\Tests\Functional\Driver\SQLServer\Query;
 
 // phpcs:ignore
+use Cycle\Database\Driver\SQLServer\Query\SQLServerInsertQuery;
+use Cycle\Database\Driver\SQLServer\Schema\SQLServerColumn;
+use Cycle\Database\Exception\BuilderException;
+use Cycle\Database\Injection\Fragment;
 use Cycle\Database\Tests\Functional\Driver\Common\Query\InsertQueryTest as CommonClass;
 
 /**
  * @group driver
  * @group driver-sqlserver
  */
-class InsertQueryTest extends CommonClass
+final class InsertQueryTest extends CommonClass
 {
     public const DRIVER = 'sqlserver';
+
+    public function testQueryInstance(): void
+    {
+        parent::testQueryInstance();
+        $this->assertInstanceOf(SQLServerInsertQuery::class, $this->database->insert());
+    }
+
+    public function testReturning(): void
+    {
+        $insert = $this->database->insert()->into('table')
+            ->columns('name', 'balance')
+            ->values('John Doe', 100)
+            ->returning('name');
+
+        $this->assertSameQuery(
+            'INSERT INTO {table} ({name}, {balance}) OUTPUT INSERTED.{name} VALUES (?,?)',
+            $insert
+        );
+    }
+
+    public function testMultipleReturning(): void
+    {
+        $insert = $this->database->insert()->into('table')
+            ->columns('name', 'balance')
+            ->values('John Doe', 100)
+            ->returning('name', 'created_at');
+
+        $this->assertSameQuery(
+            'INSERT INTO {table} ({name}, {balance}) OUTPUT INSERTED.{name}, INSERTED.{created_at} VALUES (?,?)',
+            $insert
+        );
+    }
+
+    public function testReturningWithFragment(): void
+    {
+        $insert = $this->database->insert()->into('table')
+            ->columns('name', 'balance')
+            ->values('John Doe', 100)
+            ->returning(new Fragment('INSERTED.[name] as [full_name]'));
+
+        $this->assertSameQuery(
+            'INSERT INTO {table} ({name}, {balance}) OUTPUT INSERTED.{name} as {full_name} VALUES (?,?)',
+            $insert
+        );
+    }
+
+    public function testMultipleReturningWithFragment(): void
+    {
+        $insert = $this->database->insert()->into('table')
+            ->columns('name', 'balance')
+            ->values('John Doe', 100)
+            ->returning('name', new Fragment('INSERTED.[created_at] as [date]'));
+
+        $this->assertSameQuery(
+            'INSERT INTO {table} ({name}, {balance}) OUTPUT INSERTED.{name}, INSERTED.{created_at} as {date} VALUES (?,?)',
+            $insert
+        );
+    }
+
+    public function testReturningWithDefaultValues(): void
+    {
+        $insert = $this->database->insert()->into('table')->values([])->returning('created_at');
+
+        $this->assertSameQuery(
+            'INSERT INTO {table} OUTPUT INSERTED.[created_at] DEFAULT VALUES',
+            $insert
+        );
+    }
+
+    public function testReturningValuesFromDatabase(): void
+    {
+        $schema = $this->schema('returning_values');
+        $schema->primary('id');
+        $schema->string('name');
+        $schema->datetime('created_at', defaultValue: SQLServerColumn::DATETIME_NOW);
+        $schema->save();
+
+        $returning = $this->database
+            ->insert('returning_values')
+            ->values(['name' => 'foo'])
+            ->returning('id', 'created_at')
+            ->run();
+
+        $this->assertSame(1, (int) $returning['id']);
+        $this->assertIsString($returning['created_at']);
+        $this->assertNotFalse(\strtotime($returning['created_at']));
+
+        $returning = $this->database
+            ->insert('returning_values')
+            ->values(['name' => 'foo'])
+            ->returning('id', new Fragment('INSERTED.[created_at] as [datetime]'))
+            ->run();
+
+        $this->assertSame(2, (int) $returning['id']);
+        $this->assertIsString($returning['datetime']);
+        $this->assertNotFalse(\strtotime($returning['datetime']));
+    }
+
+    public function testReturningSingleValueFromDatabase(): void
+    {
+        $schema = $this->schema('returning_value');
+        $schema->primary('id');
+        $schema->string('name');
+        $schema->datetime('created_at', defaultValue: SQLServerColumn::DATETIME_NOW);
+        $schema->save();
+
+        $returning = $this->database
+            ->insert('returning_value')
+            ->values(['name' => 'foo'])
+            ->returning(new Fragment('INSERTED.[created_at] as [datetime]'))
+            ->run();
+
+        $this->assertIsString($returning);
+        $this->assertNotFalse(\strtotime($returning));
+    }
+
+    public function testReturningValuesFromDatabaseWithDefaultValuesInsert(): void
+    {
+        $schema = $this->schema('returning_value');
+        $schema->primary('id');
+        $schema->datetime('created_at', defaultValue: SQLServerColumn::DATETIME_NOW);
+        $schema->datetime('updated_at', defaultValue: SQLServerColumn::DATETIME_NOW);
+        $schema->save();
+
+        $returning = $this->database
+            ->insert('returning_value')
+            ->values([])
+            ->returning('updated_at', new Fragment('INSERTED.[created_at] as [created]'))
+            ->run();
+
+        $this->assertIsString($returning['created']);
+        $this->assertNotFalse(\strtotime($returning['created']));
+
+        $this->assertIsString($returning['updated_at']);
+        $this->assertNotFalse(\strtotime($returning['updated_at']));
+    }
+
+    public function testCustomReturningShouldContainColumns(): void
+    {
+        $this->expectException(BuilderException::class);
+        $this->expectExceptionMessage('RETURNING clause should contain at least 1 column.');
+
+        $this->database->insert()->into('table')
+            ->columns('name', 'balance')
+            ->values('John Doe', 100)
+            ->returning();
+    }
 }

tests/Database/Unit/Driver/SQLServer/Query/SQLServerInsertQueryTest.php

@@ -0,0 +1,21 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Cycle\Database\Tests\Unit\Driver\SQLServer\Query;
+
+use Cycle\Database\Driver\DriverInterface;
+use Cycle\Database\Driver\SQLServer\Query\SQLServerInsertQuery;
+use Cycle\Database\Exception\BuilderException;
+use PHPUnit\Framework\TestCase;
+
+final class SQLServerInsertQueryTest extends TestCase
+{
+    public function testWithDriverException(): void
+    {
+        $insert = new SQLServerInsertQuery();
+
+        $this->expectException(BuilderException::class);
+        $insert->withDriver($this->createMock(DriverInterface::class));
+    }
+}