From 04df1191aeb80267726836e1552c83d9bc4cc4a9 Mon Sep 17 00:00:00 2001
From: Jonas Staudenmeir <mail@jonas-staudenmeir.de>
Date: Thu, 18 Jan 2024 00:03:57 +0100
Subject: [PATCH] Add withInitialQueryConstraint() and withQueryConstraint()
 methods

---
 README.md                                     |  26 +++--
 src/Eloquent/Traits/HasAdjacencyList.php      |  49 +-------
 src/Eloquent/Traits/HasGraphAdjacencyList.php |  49 +-------
 .../Traits/HasGraphRelationshipScopes.php     |   4 +
 src/Eloquent/Traits/HasQueryConstraints.php   | 106 ++++++++++++++++++
 .../Traits/HasRecursiveRelationshipScopes.php |   4 +
 tests/Graph/EloquentTest.php                  |  54 +++++++--
 tests/Tree/EloquentTest.php                   |  42 ++++++-
 8 files changed, 210 insertions(+), 124 deletions(-)
 create mode 100644 src/Eloquent/Traits/HasQueryConstraints.php

diff --git a/README.md b/README.md
index 955f9f6..b42eb92 100644
--- a/README.md
+++ b/README.md
@@ -48,7 +48,7 @@ Supports Laravel 5.5.29+.
 - [Path](#path)
 - [Custom Paths](#custom-paths)
 - [Nested Results](#nested-results)
-- [Recursive Query Constraints](#recursive-query-constraints)
+- [Initial & Recursive Query Constraints](#initial--recursive-query-constraints)
 - [Custom Relationships](#custom-relationships)
 - [Concatenation](#concatenation)
 
@@ -363,19 +363,22 @@ This recursively sets `children` relationships:
 ]
 ```
 
-#### Recursive Query Constraints
+#### Initial & Recursive Query Constraints
 
-You can add custom constraints to the CTE's recursive query. Consider a query where you want to traverse a tree while
-skipping inactive users and their subtrees:
+You can add custom constraints to the CTE's initial and recursive query. Consider a query where you want to traverse a
+tree while skipping inactive users and their descendants:
 
  ```php
-$tree = User::withRecursiveQueryConstraint(function (Builder $query) {
+$tree = User::withQueryConstraint(function (Builder $query) {
     $query->where('users.active', true);
 }, function () {
     return User::tree()->get();
 });
  ```
 
+ You can also add a custom constraint to only the initial or recursive query using `withInitialQueryConstraint()`/
+ `withRecursiveQueryConstraint()`.
+
 #### Custom Relationships
 
 You can also define custom relationships to retrieve related models recursively.
@@ -663,7 +666,7 @@ Supports Laravel 9+.
 - [Path](#graphs-path)
 - [Custom Paths](#graphs-custom-paths)
 - [Nested Results](#graphs-nested-results)
-- [Recursive Query Constraints](#graphs-recursive-query-constraints)
+- [Initial & Recursive Query Constraints](#graphs-initial--recursive-query-constraints)
 
 #### <a name="graphs-getting-started">Getting Started</a>
 
@@ -1026,19 +1029,22 @@ This recursively sets `children` relationships:
 ]
 ```
 
-#### <a name="graphs-recursive-query-constraints">Recursive Query Constraints</a>
+#### <a name="graphs-initial--recursive-query-constraints">Initial & Recursive Query Constraints</a>
 
-You can add custom constraints to the CTE's recursive query. Consider a query where you want to traverse a node's
-descendants while skipping inactive nodes and their subgraphs:
+You can add custom constraints to the CTE's initial and recursive query. Consider a query where you want to traverse a
+node's descendants while skipping inactive nodes and their descendants:
 
  ```php
-$descendants = Node::withRecursiveQueryConstraint(function (Builder $query) {
+$descendants = Node::withQueryConstraint(function (Builder $query) {
     $query->where('nodes.active', true);
 }, function () {
     return Node::find($id)->descendants;
 });
  ```
 
+You can also add a custom constraint to only the initial or recursive query using `withInitialQueryConstraint()`/
+`withRecursiveQueryConstraint()`.
+
 ### Package Conflicts
 
 - `staudenmeir/eloquent-eager-limit`: Replace both packages
diff --git a/src/Eloquent/Traits/HasAdjacencyList.php b/src/Eloquent/Traits/HasAdjacencyList.php
index c2c35c7..75019c2 100644
--- a/src/Eloquent/Traits/HasAdjacencyList.php
+++ b/src/Eloquent/Traits/HasAdjacencyList.php
@@ -15,15 +15,9 @@
 trait HasAdjacencyList
 {
     use HasOfDescendantsRelationships;
+    use HasQueryConstraints;
     use HasRecursiveRelationshipScopes;
 
-    /**
-     * The additional constraint for the recursive query.
-     *
-     * @var callable|null
-     */
-    public static $recursiveQueryConstraint;
-
     /**
      * Get the name of the parent key column.
      *
@@ -414,45 +408,4 @@ public function newCollection(array $models = [])
     {
         return new Collection($models);
     }
-
-    /**
-     * Set an additional constraint for the recursive query.
-     *
-     * @param callable $constraint
-     * @return void
-     */
-    public static function setRecursiveQueryConstraint(callable $constraint)
-    {
-        static::$recursiveQueryConstraint = $constraint;
-    }
-
-    /**
-     * Unset the additional constraint for the recursive query.
-     *
-     * @return void
-     */
-    public static function unsetRecursiveQueryConstraint()
-    {
-        static::$recursiveQueryConstraint = null;
-    }
-
-    /**
-     * Execute a query with an additional constraint for the recursive query.
-     *
-     * @param callable $constraint
-     * @param callable $query
-     * @return mixed
-     */
-    public static function withRecursiveQueryConstraint(callable $constraint, callable $query)
-    {
-        $previous = static::$recursiveQueryConstraint;
-
-        static::$recursiveQueryConstraint = $constraint;
-
-        $result = $query();
-
-        static::$recursiveQueryConstraint = $previous;
-
-        return $result;
-    }
 }
diff --git a/src/Eloquent/Traits/HasGraphAdjacencyList.php b/src/Eloquent/Traits/HasGraphAdjacencyList.php
index b1042fe..ffbc98a 100644
--- a/src/Eloquent/Traits/HasGraphAdjacencyList.php
+++ b/src/Eloquent/Traits/HasGraphAdjacencyList.php
@@ -12,13 +12,7 @@
 trait HasGraphAdjacencyList
 {
     use HasGraphRelationshipScopes;
-
-    /**
-     * The additional constraint for the recursive query.
-     *
-     * @var callable|null
-     */
-    public static $recursiveQueryConstraint;
+    use HasQueryConstraints;
 
     /**
      * Get the name of the pivot table.
@@ -433,47 +427,6 @@ public function newCollection(array $models = [])
         return new Collection($models);
     }
 
-    /**
-     * Set an additional constraint for the recursive query.
-     *
-     * @param callable $constraint
-     * @return void
-     */
-    public static function setRecursiveQueryConstraint(callable $constraint): void
-    {
-        static::$recursiveQueryConstraint = $constraint;
-    }
-
-    /**
-     * Unset the additional constraint for the recursive query.
-     *
-     * @return void
-     */
-    public static function unsetRecursiveQueryConstraint(): void
-    {
-        static::$recursiveQueryConstraint = null;
-    }
-
-    /**
-     * Execute a query with an additional constraint for the recursive query.
-     *
-     * @param callable $constraint
-     * @param callable $query
-     * @return mixed
-     */
-    public static function withRecursiveQueryConstraint(callable $constraint, callable $query): mixed
-    {
-        $previous = static::$recursiveQueryConstraint;
-
-        static::$recursiveQueryConstraint = $constraint;
-
-        $result = $query();
-
-        static::$recursiveQueryConstraint = $previous;
-
-        return $result;
-    }
-
     /**
      * Execute a query with a maximum depth constraint for the recursive query.
      *
diff --git a/src/Eloquent/Traits/HasGraphRelationshipScopes.php b/src/Eloquent/Traits/HasGraphRelationshipScopes.php
index 84a0d16..6fa5afe 100644
--- a/src/Eloquent/Traits/HasGraphRelationshipScopes.php
+++ b/src/Eloquent/Traits/HasGraphRelationshipScopes.php
@@ -142,6 +142,10 @@ protected function getInitialQuery(
 
         $constraint($query);
 
+        if (static::$initialQueryConstraint) {
+            (static::$initialQueryConstraint)($query);
+        }
+
         return $query;
     }
 
diff --git a/src/Eloquent/Traits/HasQueryConstraints.php b/src/Eloquent/Traits/HasQueryConstraints.php
new file mode 100644
index 0000000..b13a5cb
--- /dev/null
+++ b/src/Eloquent/Traits/HasQueryConstraints.php
@@ -0,0 +1,106 @@
+<?php
+
+namespace Staudenmeir\LaravelAdjacencyList\Eloquent\Traits;
+
+trait HasQueryConstraints
+{
+    /**
+     * The additional constraint for the initial query.
+     *
+     * @var callable|null
+     */
+    public static $initialQueryConstraint;
+
+    /**
+     * The additional constraint for the recursive query.
+     *
+     * @var callable|null
+     */
+    public static $recursiveQueryConstraint;
+
+    /**
+     * Execute a query with an additional constraint for the initial query.
+     *
+     * @param callable $constraint
+     * @param callable $query
+     * @return mixed
+     */
+    public static function withInitialQueryConstraint(callable $constraint, callable $query): mixed
+    {
+        $previous = static::$initialQueryConstraint;
+
+        static::$initialQueryConstraint = $constraint;
+
+        $result = $query();
+
+        static::$initialQueryConstraint = $previous;
+
+        return $result;
+    }
+
+    /**
+     * Execute a query with an additional constraint for the recursive query.
+     *
+     * @param callable $constraint
+     * @param callable $query
+     * @return mixed
+     */
+    public static function withRecursiveQueryConstraint(callable $constraint, callable $query): mixed
+    {
+        $previous = static::$recursiveQueryConstraint;
+
+        static::$recursiveQueryConstraint = $constraint;
+
+        $result = $query();
+
+        static::$recursiveQueryConstraint = $previous;
+
+        return $result;
+    }
+
+    /**
+     * Set an additional constraint for the recursive query.
+     *
+     * @param callable $constraint
+     * @return void
+     */
+    public static function setRecursiveQueryConstraint(callable $constraint): void
+    {
+        static::$recursiveQueryConstraint = $constraint;
+    }
+
+    /**
+     * Unset the additional constraint for the recursive query.
+     *
+     * @return void
+     */
+    public static function unsetRecursiveQueryConstraint(): void
+    {
+        static::$recursiveQueryConstraint = null;
+    }
+
+    /**
+     * Execute a query with an additional constraint for the initial and recursive query.
+     *
+     * @param callable $constraint
+     * @param callable $query
+     * @return mixed
+     */
+    public static function withQueryConstraint(
+        callable $constraint,
+        callable $query
+    ): mixed {
+        $previousInitialConstraint = static::$initialQueryConstraint;
+        $previousRecursiveConstraint = static::$recursiveQueryConstraint;
+
+        static::$initialQueryConstraint = $constraint;
+        static::$recursiveQueryConstraint = $constraint;
+
+        $result = $query();
+
+        static::$initialQueryConstraint = $previousInitialConstraint;
+        static::$recursiveQueryConstraint = $previousRecursiveConstraint;
+
+        return $result;
+    }
+}
diff --git a/src/Eloquent/Traits/HasRecursiveRelationshipScopes.php b/src/Eloquent/Traits/HasRecursiveRelationshipScopes.php
index aa4f665..0c71028 100644
--- a/src/Eloquent/Traits/HasRecursiveRelationshipScopes.php
+++ b/src/Eloquent/Traits/HasRecursiveRelationshipScopes.php
@@ -189,6 +189,10 @@ protected function getInitialQuery(ExpressionGrammar $grammar, callable $constra
 
         $constraint($query);
 
+        if (static::$initialQueryConstraint) {
+            (static::$initialQueryConstraint)($query);
+        }
+
         return $query;
     }
 
diff --git a/tests/Graph/EloquentTest.php b/tests/Graph/EloquentTest.php
index 40c2bbc..232b30e 100644
--- a/tests/Graph/EloquentTest.php
+++ b/tests/Graph/EloquentTest.php
@@ -109,36 +109,66 @@ public function testScopeDepthFirst()
         $this->assertEquals([2, 5, 7, 8, 8, 3, 6, 4, 5, 7, 8, 8], $nodes->pluck('id')->all());
     }
 
+    public function testWithInitialQueryConstraint()
+    {
+        $nodes = Node::withInitialQueryConstraint(function (Builder $query) {
+            $query->where('edges.weight', '<', 2);
+        }, function () {
+            return Node::find(1)->descendants()->orderBy('id')->get();
+        });
+
+        $this->assertEquals([2, 5, 7, 8, 8], $nodes->pluck('id')->all());
+
+        $nodes = Node::find(1)->descendants()->orderBy('id')->get();
+
+        $this->assertEquals([2, 3, 4, 5, 5, 6, 7, 7, 8, 8, 8, 8], $nodes->pluck('id')->all());
+    }
+
+    public function testWithRecursiveQueryConstraint()
+    {
+        $nodes = Node::withRecursiveQueryConstraint(function (Builder $query) {
+            $query->where('edges.weight', '<', 5);
+        }, function () {
+            return Node::find(1)->descendants()->orderBy('id')->get();
+        });
+
+        $this->assertEquals([2, 3, 4, 5], $nodes->pluck('id')->all());
+
+        $nodes = Node::find(1)->descendants()->orderBy('id')->get();
+
+        $this->assertEquals([2, 3, 4, 5, 5, 6, 7, 7, 8, 8, 8, 8], $nodes->pluck('id')->all());
+    }
+
     public function testSetRecursiveQueryConstraint()
     {
         Node::setRecursiveQueryConstraint(function (Builder $query) {
-            $query->where('pivot_weight', '<', 7);
+            $query->where('edges.weight', '<', 5);
         });
 
-        $nodes = Node::find(2)->descendants()->orderBy('id')->get();
+        $nodes = Node::find(1)->descendants()->orderBy('id')->get();
 
-        $this->assertEquals([5, 7, 8], $nodes->pluck('id')->all());
+        $this->assertEquals([2, 3, 4, 5], $nodes->pluck('id')->all());
 
         Node::unsetRecursiveQueryConstraint();
 
-        $nodes = Node::find(2)->descendants()->orderBy('id')->get();
+        $nodes = Node::find(1)->descendants()->orderBy('id')->get();
 
-        $this->assertEquals([5, 7, 8, 8], $nodes->pluck('id')->all());
+        $this->assertEquals([2, 3, 4, 5, 5, 6, 7, 7, 8, 8, 8, 8], $nodes->pluck('id')->all());
     }
 
-    public function testWithRecursiveQueryConstraint()
+    public function testWithQueryConstraint()
     {
-        $nodes = Node::withRecursiveQueryConstraint(function (Builder $query) {
-            $query->where('pivot_weight', '<', 7);
+        $nodes = Node::withQueryConstraint(function (Builder $query) {
+            $query->where('edges.weight', '<', 4);
         }, function () {
-            return Node::find(2)->descendants()->orderBy('id')->get();
+            return Node::find(1)->descendants()->orderBy('id')->get();
         });
 
-        $this->assertEquals([5, 7, 8], $nodes->pluck('id')->all());
+        $this->assertEquals([2, 3, 4], $nodes->pluck('id')->all());
 
-        $nodes = Node::find(2)->descendants()->orderBy('id')->get();
+        $nodes = Node::find(1)->descendants()->orderBy('id')->get();
 
-        $this->assertEquals([5, 7, 8, 8], $nodes->pluck('id')->all());
+        $this->assertEquals([2, 3, 4, 5, 5, 6, 7, 7, 8, 8, 8, 8], $nodes->pluck('id')->all());
     }
 
     public function testWithMaxDepth()
diff --git a/tests/Tree/EloquentTest.php b/tests/Tree/EloquentTest.php
index 6e2a29c..1da6b8b 100644
--- a/tests/Tree/EloquentTest.php
+++ b/tests/Tree/EloquentTest.php
@@ -162,15 +162,45 @@ public function testScopeDepthFirstWithStringKey()
         $this->assertEquals(['a', 'b', 'c', 'd'], $categories->pluck('id')->all());
     }
 
+    public function testWithInitialQueryConstraint()
+    {
+        $users = User::withInitialQueryConstraint(function (Builder $query) {
+            $query->where('users.id', '<>', 1);
+        }, function () {
+            return User::tree()->orderBy('id')->get();
+        });
+
+        $this->assertEquals([11, 12], $users->pluck('id')->all());
+
+        $users = User::tree()->orderBy('id')->get();
+
+        $this->assertEquals([1, 2, 3, 4, 5, 6, 7, 8, 9, 11, 12], $users->pluck('id')->all());
+    }
+
+    public function testWithRecursiveQueryConstraint()
+    {
+        $users = User::withRecursiveQueryConstraint(function (Builder $query) {
+            $query->where('users.id', '<', 5);
+        }, function () {
+            return User::tree()->orderBy('id')->get();
+        });
+
+        $this->assertEquals([1, 2, 3, 4, 11], $users->pluck('id')->all());
+
+        $users = User::tree()->orderBy('id')->get();
+
+        $this->assertEquals([1, 2, 3, 4, 5, 6, 7, 8, 9, 11, 12], $users->pluck('id')->all());
+    }
+
     public function testSetRecursiveQueryConstraint()
     {
         User::setRecursiveQueryConstraint(function (Builder $query) {
-            $query->where('users.parent_id', '<', 4);
+            $query->where('users.id', '<', 5);
         });
 
         $users = User::tree()->orderBy('id')->get();
 
-        $this->assertEquals([1, 2, 3, 4, 5, 6, 11], $users->pluck('id')->all());
+        $this->assertEquals([1, 2, 3, 4, 11], $users->pluck('id')->all());
 
         User::unsetRecursiveQueryConstraint();
 
@@ -179,15 +209,15 @@ public function testSetRecursiveQueryConstraint()
         $this->assertEquals([1, 2, 3, 4, 5, 6, 7, 8, 9, 11, 12], $users->pluck('id')->all());
     }
 
-    public function testWithRecursiveQueryConstraint()
+    public function testWithQueryConstraint()
     {
-        $users = User::withRecursiveQueryConstraint(function (Builder $query) {
-            $query->where('users.parent_id', '<', 4);
+        $users = User::withQueryConstraint(function (Builder $query) {
+            $query->where('users.id', '<', 5);
         }, function () {
             return User::tree()->orderBy('id')->get();
         });
 
-        $this->assertEquals([1, 2, 3, 4, 5, 6, 11], $users->pluck('id')->all());
+        $this->assertEquals([1, 2, 3, 4], $users->pluck('id')->all());
 
         $users = User::tree()->orderBy('id')->get();