Add support for PARTITION BY clause in window function

src/Query/Mysql/Over.php

@@ -2,29 +2,48 @@
 
 namespace DoctrineExtensions\Query\Mysql;
 
+use Closure;
 use Doctrine\ORM\Query\AST\ArithmeticExpression;
 use Doctrine\ORM\Query\AST\Functions\FunctionNode;
 use Doctrine\ORM\Query\AST\OrderByClause;
 use Doctrine\ORM\Query\Parser;
 use Doctrine\ORM\Query\SqlWalker;
 use Doctrine\ORM\Query\TokenType;
 
+use function array_map;
 use function count;
+use function implode;
+use function strtolower;
 use function trim;
 
 class Over extends FunctionNode
 {
     /** @var ArithmeticExpression */
-    private $arithmeticExpression;
+    private $windowFunction;
+
+    /** @var ArithmeticExpression[] */
+    private $partitionByClause;
 
     /** @var OrderByClause|null */
     private $orderByClause;
 
     public function getSql(SqlWalker $sqlWalker): string
     {
-        return isset($this->orderByClause) && count($this->orderByClause->orderByItems) > 0
-            ? $sqlWalker->walkArithmeticExpression($this->arithmeticExpression) . ' OVER (' . trim($sqlWalker->walkOrderByClause($this->orderByClause)) . ')'
-            : $sqlWalker->walkArithmeticExpression($this->arithmeticExpression) . ' OVER ()';
+        $sql           = $sqlWalker->walkArithmeticExpression($this->windowFunction) . ' OVER (';
+        $overClauseSQL = [];
+        if ($this->partitionByClause !== null && count($this->partitionByClause) !== 0) {
+            $partitionByClauseSQL = array_map(
+                Closure::fromCallable([$sqlWalker, 'walkArithmeticExpression']),
+                $this->partitionByClause
+            );
+            $overClauseSQL[]      = 'PARTITION BY ' . trim(implode(', ', $partitionByClauseSQL));
+        }
+
+        if ($this->orderByClause !== null && count($this->orderByClause->orderByItems) !== 0) {
+            $overClauseSQL[] = trim($sqlWalker->walkOrderByClause($this->orderByClause));
+        }
+
+        return $sql . implode(' ', $overClauseSQL) . ')';
     }
 
     public function parse(Parser $parser): void
@@ -33,10 +52,33 @@ public function parse(Parser $parser): void
 
         $parser->match(TokenType::T_IDENTIFIER);
         $parser->match(TokenType::T_OPEN_PARENTHESIS);
-        $this->arithmeticExpression = $parser->ArithmeticExpression();
-        if (! $lexer->isNextToken(TokenType::T_CLOSE_PARENTHESIS)) {
+        $this->windowFunction = $parser->ArithmeticExpression();
+        if (
+            ! $lexer->isNextToken(TokenType::T_CLOSE_PARENTHESIS)
+            && $lexer->isNextToken(TokenType::T_COMMA)
+        ) {
             $parser->match(TokenType::T_COMMA);
-            $this->orderByClause = $parser->OrderByClause();
+            if (
+                $lexer->isNextToken(TokenType::T_IDENTIFIER)
+                && strtolower($lexer->lookahead->value) === 'partition'
+            ) {
+                $parser->match(TokenType::T_IDENTIFIER);
+                $parser->match(TokenType::T_BY);
+
+                $this->partitionByClause[] = $parser->ArithmeticExpression();
+
+                while ($lexer->isNextToken(TokenType::T_COMMA)) {
+                    $parser->match(TokenType::T_COMMA);
+                    $this->partitionByClause[] = $parser->ArithmeticExpression();
+                }
+            }
+
+            if (
+                $lexer->isNextToken(TokenType::T_ORDER)
+                && strtolower($lexer->lookahead->value) === 'order'
+            ) {
+                $this->orderByClause = $parser->OrderByClause();
+            }
         }
 
         $parser->match(TokenType::T_CLOSE_PARENTHESIS);

tests/Query/Mysql/LagTest.php

@@ -27,20 +27,5 @@ public function testLag(): void
             'SELECT LAG(COUNT(b.id), 5, 5 + 5) from DoctrineExtensions\Tests\Entities\Blank b',
             'SELECT LAG(COUNT(b0_.id), 5, 5 + 5) AS sclr_0 FROM Blank b0_'
         );
-
-        $this->assertDqlProducesSql(
-            'SELECT OVER(LAG(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
-            'SELECT LAG(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
-        );
-
-        $this->assertDqlProducesSql(
-            'SELECT OVER(LAG(COUNT(b.id))) from DoctrineExtensions\Tests\Entities\Blank b',
-            'SELECT LAG(COUNT(b0_.id)) OVER () AS sclr_0 FROM Blank b0_'
-        );
-
-        $this->assertDqlProducesSql(
-            'SELECT OVER(COUNT(b.id) - LAG(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
-            'SELECT COUNT(b0_.id) - LAG(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
-        );
     }
 }

tests/Query/Mysql/LeadTest.php

@@ -28,19 +28,6 @@ public function testLag(): void
             'SELECT LEAD(COUNT(b0_.id), 5, 5 + 5) AS sclr_0 FROM Blank b0_'
         );
 
-        $this->assertDqlProducesSql(
-            'SELECT OVER(LEAD(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
-            'SELECT LEAD(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
-        );
 
-        $this->assertDqlProducesSql(
-            'SELECT OVER(LEAD(COUNT(b.id))) from DoctrineExtensions\Tests\Entities\Blank b',
-            'SELECT LEAD(COUNT(b0_.id)) OVER () AS sclr_0 FROM Blank b0_'
-        );
-
-        $this->assertDqlProducesSql(
-            'SELECT OVER(COUNT(b.id) - LEAD(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
-            'SELECT COUNT(b0_.id) - LEAD(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
-        );
     }
 }

tests/Query/Mysql/OverTest.php

@@ -0,0 +1,57 @@
+<?php
+
+namespace DoctrineExtensions\Tests\Query\Mysql;
+
+use DoctrineExtensions\Tests\Query\MysqlTestCase;
+
+class OverTest extends MysqlTestCase
+{
+    public function testPartitionBy(): void
+    {
+        $this->assertDqlProducesSql(
+            'SELECT OVER(LEAD(COUNT(b.id)), PARTITION BY b.id ORDER BY b.id DESC) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT LEAD(COUNT(b0_.id)) OVER (PARTITION BY b0_.id ORDER BY b0_.id DESC) AS sclr_0 FROM Blank b0_'
+        );
+
+        $this->assertDqlProducesSql(
+            'SELECT OVER(LEAD(COUNT(b.id)), PARTITION BY COUNT(b.id) ORDER BY COUNT(b.id) DESC) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT LEAD(COUNT(b0_.id)) OVER (PARTITION BY COUNT(b0_.id) ORDER BY COUNT(b0_.id) DESC) AS sclr_0 FROM Blank b0_'
+        );
+    }
+
+    public function testLead(): void
+    {
+        $this->assertDqlProducesSql(
+            'SELECT OVER(LEAD(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT LEAD(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
+        );
+
+        $this->assertDqlProducesSql(
+            'SELECT OVER(LEAD(COUNT(b.id))) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT LEAD(COUNT(b0_.id)) OVER () AS sclr_0 FROM Blank b0_'
+        );
+
+        $this->assertDqlProducesSql(
+            'SELECT OVER(COUNT(b.id) - LEAD(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT COUNT(b0_.id) - LEAD(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
+        );
+    }
+
+    public function testLag(): void
+    {
+        $this->assertDqlProducesSql(
+            'SELECT OVER(LAG(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT LAG(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
+        );
+
+        $this->assertDqlProducesSql(
+            'SELECT OVER(LAG(COUNT(b.id))) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT LAG(COUNT(b0_.id)) OVER () AS sclr_0 FROM Blank b0_'
+        );
+
+        $this->assertDqlProducesSql(
+            'SELECT OVER(COUNT(b.id) - LAG(COUNT(b.id)), ORDER BY b.id) from DoctrineExtensions\Tests\Entities\Blank b',
+            'SELECT COUNT(b0_.id) - LAG(COUNT(b0_.id)) OVER (ORDER BY b0_.id ASC) AS sclr_0 FROM Blank b0_'
+        );
+    }
+}