Merge pull request #22 from xsuchy09/feature/direct-pit-and-search-after-support

Direct support for pit and search_after in QueryBuilder

Author: erichard

src/QueryBuilder.php

@@ -21,6 +21,10 @@ class QueryBuilder
 
     private ?int $size = null;
 
+    private string|array|null $pit = null;
+
+    private ?array $searchAfter = null;
+
     private ?QueryInterface $query = null;
 
     private ?array $fields = null;
@@ -59,6 +63,24 @@ public function setSize(int $size): self
         return $this;
     }
 
+    public function setPit(string|array|null $pit): self
+    {
+        if (is_array($pit) && !isset($pit['id'])) {
+            throw new \InvalidArgumentException('Parameter "pit" is not valid. Value with key "id" is not set.');
+        }
+
+        $this->pit = $pit;
+
+        return $this;
+    }
+
+    public function setSearchAfter(?array $searchAfter): self
+    {
+        $this->searchAfter = $searchAfter;
+
+        return $this;
+    }
+
     public function setQuery(QueryInterface $query): self
     {
         $this->query = $query;
@@ -113,6 +135,18 @@ public function build(): array
             $query['size'] = $this->size;
         }
 
+        if (null !== $this->pit) {
+            if (is_string($this->pit)) {
+                $query['pit'] = ['id' => $this->pit];
+            } else {
+                $query['pit'] = $this->pit;
+            }
+        }
+
+        if (null !== $this->searchAfter) {
+            $query['search_after'] = $this->searchAfter;
+        }
+
         if (null !== $this->source) {
             $query['_source'] = $this->source;
         }
@@ -155,6 +189,16 @@ public function getSize(): ?int
         return $this->size;
     }
 
+    public function getPit(): string|array|null
+    {
+        return $this->pit;
+    }
+
+    public function getSearchAfter(): ?array
+    {
+        return $this->searchAfter;
+    }
+
     public function getQuery(): ?QueryInterface
     {
         return $this->query;

tests/QueryBuilderTest.php

@@ -64,6 +64,28 @@ public function testItSetTheFrom(): void
         $this->assertEquals(50, $query['from']);
     }
 
+    public function testItSetThePitAsString(): void
+    {
+        $queryBuilder = new QueryBuilder();
+
+        $queryBuilder->setPit('pit-as-string');
+
+        $query = $queryBuilder->build();
+
+        $this->assertEquals(['id' => 'pit-as-string'], $query['pit']);
+    }
+
+    public function testItSetThePitAsArray(): void
+    {
+        $queryBuilder = new QueryBuilder();
+
+        $queryBuilder->setPit(['id' => 'pit-as-array', 'keep_alive' => '1m']);
+
+        $query = $queryBuilder->build();
+
+        $this->assertEquals(['id' => 'pit-as-array', 'keep_alive' => '1m'], $query['pit']);
+    }
+
     public function testItAllowToSort(): void
     {
         $queryBuilder = new QueryBuilder();