Correct comment functionality, add comment constraint exception

src/Exceptions/RLSCommentConstraintException.php

@@ -0,0 +1,15 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Stancl\Tenancy\Exceptions;
+
+use Exception;
+
+class RLSCommentConstraintException extends Exception
+{
+    public function __construct(string|null $message = null)
+    {
+        parent::__construct($message ?? "Invalid comment constraint.");
+    }
+}

src/RLS/PolicyManagers/TableRLSManager.php

@@ -7,6 +7,7 @@ namespace Stancl\Tenancy\RLS\PolicyManagers;
 use Illuminate\Database\DatabaseManager;
 use Illuminate\Support\Str;
 use Stancl\Tenancy\Database\Exceptions\RecursiveRelationshipException;
+use Stancl\Tenancy\Exceptions\RLSCommentConstraintException;
 
 // todo@samuel logical + structural refactor. the tree generation could use some dynamic programming optimizations
 class TableRLSManager implements RLSPolicyManager
@@ -81,7 +82,7 @@ class TableRLSManager implements RLSPolicyManager
             $table = str($table)->afterLast('.')->toString();
 
             // For each table, we get a list of all foreign key columns
-            $foreignKeys = collect($builder->getForeignKeys($table))->merge($this->getFakeForeignKeys($table))->map(function ($foreign) use ($table) {
+            $foreignKeys = collect($builder->getForeignKeys($table))->merge($this->getCommentConstraints($table))->map(function ($foreign) use ($table) {
                 return $this->formatForeignKey($foreign, $table);
             });
 
@@ -111,7 +112,7 @@ class TableRLSManager implements RLSPolicyManager
     {
         // If the foreign key has a comment of 'no-rls', we skip it
         // Also skip the foreign key if implicit scoping is off and the foreign key has no comment
-        if ($foreign['comment'] === 'no-rls' || (! static::$scopeByDefault && $foreign['comment'] === null)) {
+        if ($this->shouldSkipPathLeadingThrough($foreign)) {
             return;
         }
 
@@ -125,48 +126,79 @@ class TableRLSManager implements RLSPolicyManager
             $paths[] = $currentPath;
         } else {
             // If not, recursively generate paths for the foreign table
-            foreach (array_merge($this->database->getSchemaBuilder()->getForeignKeys($foreign['foreignTable'], $this->getFakeForeignKeys($table))) as $nextConstraint) {
+            foreach (array_merge(
+                $this->database->getSchemaBuilder()->getForeignKeys($foreign['foreignTable']),
+                $this->getCommentConstraints($foreign['foreignTable'])
+            ) as $nextConstraint) {
                 $this->generatePaths($table, $this->formatForeignKey($nextConstraint, $foreign['foreignTable']), $paths, $currentPath);
             }
         }
     }
 
-    protected function getFakeForeignKeys(string $tableName): array
+    protected function shouldSkipPathLeadingThrough(array $foreignKey): bool
+    {
+        // If the foreign key has a comment of 'no-rls', we skip it
+        // Also skip the foreign key if implicit scoping is off and the foreign key has no comment
+        $pathExplicitlySkipped = $foreignKey['comment'] === 'no-rls';
+        $pathImplicitlySkipped = ! static::$scopeByDefault && (
+            ! isset($foreignKey['comment']) ||
+            (is_string($foreignKey['comment']) && ! Str::startsWith($foreignKey['comment'], 'rls'))
+        );
+
+        return $pathExplicitlySkipped || $pathImplicitlySkipped;
+    }
+
+    /**
+     * Retrieve comment-based constraints for a table. These are columns with comments in the format:
+     *   "rls <foreign_table>.<foreign_column>"
+     *
+     * Throws an exception if the comment is formatted incorrectly or if the referenced table/column does not exist.
+     */
+    protected function getCommentConstraints(string $tableName): array
     {
         $columns = $this->database->getSchemaBuilder()->getColumns($tableName);
+        $schemaBuilder = $this->database->getSchemaBuilder();
 
-        $fakeForeignKeys = array_filter($columns, function ($column) {
+        $commentConstraints = array_filter($columns, function ($column) {
-            // Constraint comment should be "rls <foreign_table>.<foreign_column>"
+            return (isset($column['comment']) && is_string($column['comment']))
-            if ($column['comment']) {
+                && Str::startsWith($column['comment'], 'rls ');
-                return Str::contains($column['comment'], 'rls ');
-            }
-
-            return false;
         });
 
-        if ($tableName === 'non_constrained_posts') {
+        return array_map(function ($commentConstraint) use ($schemaBuilder, $tableName) {
-            dump(array_map(function ($fakeForeignKey) {
+            $comment = $commentConstraint['comment'];
-                // Constraint comment is "rls <foreign_table>.<foreign_column>"
+            $constraintString = Str::after($comment, 'rls ');
-                $constraint = explode('.', Str::after($fakeForeignKey['comment'], 'rls '));
+            $constraint = explode('.', $constraintString);
 
-                return [
+            // Validate comment constraint format
-                    'foreign_table' => $constraint[0],
+            if (count($constraint) !== 2 || empty($constraint[0]) || empty($constraint[1])) {
-                    'foreign_columns' => [$constraint[1]],
+                throw new RLSCommentConstraintException("Incorrectly formatted comment constraint on {$tableName}.{$commentConstraint['name']}: '{$comment}'");
-                    'columns' => [$fakeForeignKey['name']],
+            }
-                ];
-            }, $fakeForeignKeys));
-        }
 
-        return array_map(function ($fakeForeignKey) {
+            $foreignTable = $constraint[0];
-            // Constraint comment is "rls <foreign_table>.<foreign_column>"
+            $foreignColumn = $constraint[1];
-            $constraint = explode('.', Str::after($fakeForeignKey['comment'], 'rls '));
+
+            // Validate table existence
+            $allTables = array_map(function ($table) {
+                return str($table)->afterLast('.')->toString();
+            }, $schemaBuilder->getTableListing(schema: $this->database->getConfig('search_path')));
+
+            if (! in_array($foreignTable, $allTables, true)) {
+                throw new RLSCommentConstraintException("Comment constraint on {$tableName}.{$commentConstraint['name']} references non-existent table '{$foreignTable}'");
+            }
+
+            // Validate column existence
+            $foreignColumns = $schemaBuilder->getColumns($foreignTable);
+            $foreignColumnNames = array_column($foreignColumns, 'name');
+            if (! in_array($foreignColumn, $foreignColumnNames, true)) {
+                throw new RLSCommentConstraintException("Comment constraint on {$tableName}.{$commentConstraint['name']} references non-existent column '{$foreignTable}.{$foreignColumn}'");
+            }
 
             return [
-                'foreign_table' => $constraint[0],
+                'foreign_table' => $foreignTable,
-                'foreign_columns' => [$constraint[1]],
+                'foreign_columns' => [$foreignColumn],
-                'columns' => [$fakeForeignKey['name']],
+                'columns' => [$commentConstraint['name']],
             ];
-        }, $fakeForeignKeys);
+        }, $commentConstraints);
     }
 
     /** Get tree's non-nullable paths. */
@@ -215,7 +247,7 @@ class TableRLSManager implements RLSPolicyManager
     }
 
     /**
-     * Formats the foreign key array retrieved by Postgres to a more readable format.
+     * Formats the foreign key retrieved by Postgres or comment-based constraint to a more readable format.
      *
      * Also provides information about whether the foreign key is nullable,
      * and the foreign key column comment. These additional details are removed

tests/RLS/TableManagerTest.php

@@ -664,6 +664,12 @@ test('table manager ignores recursive relationship if the foreign key responsibl
 });
 
 test('table manager can generate paths leading through non-constrained foreign keys', function() {
+    // Drop extra tables
+    Schema::dropIfExists('reactions');
+    Schema::dropIfExists('comments');
+    Schema::dropIfExists('posts');
+    Schema::dropIfExists('authors');
+
     Schema::create('non_constrained_users', function (Blueprint $table) {
         $table->id();
         $table->string('tenant_id')->comment('rls tenants.id'); // "fake" constraint
@@ -678,58 +684,6 @@ test('table manager can generate paths leading through non-constrained foreign k
     $manager = app(TableRLSManager::class);
 
     $expectedTrees = [
-        'authors' => [
-            // Directly related to tenants
-            'tenant_id' => [
-                [
-                    [
-                        'foreignKey' => 'tenant_id',
-                        'foreignTable' => 'tenants',
-                        'foreignId' => 'id',
-                        'nullable' => false,
-                    ]
-                ],
-            ],
-        ],
-        'comments' => [
-            // Tree starting from the post_id foreign key
-            'post_id' => [
-                [
-                    [
-                        'foreignKey' => 'post_id',
-                        'foreignTable' => 'posts',
-                        'foreignId' => 'id',
-                        'nullable' => false,
-                    ],
-                    [
-                        'foreignKey' => 'author_id',
-                        'foreignTable' => 'authors',
-                        'foreignId' => 'id',
-                        'nullable' => false,
-                    ],
-                    [
-                        'foreignKey' => 'tenant_id',
-                        'foreignTable' => 'tenants',
-                        'foreignId' => 'id',
-                        'nullable' => false,
-                    ],
-                ],
-                [
-                    [
-                        'foreignKey' => 'post_id',
-                        'foreignTable' => 'posts',
-                        'foreignId' => 'id',
-                        'nullable' => false,
-                    ],
-                    [
-                        'foreignKey' => 'tenant_id',
-                        'foreignTable' => 'tenants',
-                        'foreignId' => 'id',
-                        'nullable' => true,
-                    ],
-                ],
-            ],
-        ],
         'non_constrained_posts' => [
             'author_id' => [
                 [
@@ -762,34 +716,6 @@ test('table manager can generate paths leading through non-constrained foreign k
                 ]
             ],
         ],
-        'posts' => [
-            'author_id' => [
-                [
-                    [
-                        'foreignKey' => 'author_id',
-                        'foreignTable' => 'authors',
-                        'foreignId' => 'id',
-                        'nullable' => false,
-                    ],
-                    [
-                        'foreignKey' => 'tenant_id',
-                        'foreignTable' => 'tenants',
-                        'foreignId' => 'id',
-                        'nullable' => false,
-                    ]
-                ],
-            ],
-            'tenant_id' => [
-                [
-                    [
-                        'foreignKey' => 'tenant_id',
-                        'foreignTable' => 'tenants',
-                        'foreignId' => 'id',
-                        'nullable' => true,
-                    ]
-                ]
-            ],
-        ],
     ];
 
     expect($manager->generateTrees())->toEqual($expectedTrees);