Support wrapping a RecordingReflector in LazyQueryReflector (#607)

staabm · Jun 27, 2023 · 1bd0de3 · 1bd0de3
1 parent 9058bc9
commit 1bd0de3
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/src/QueryReflection/LazyQueryReflector.php b/src/QueryReflection/LazyQueryReflector.php
@@ -7,7 +7,7 @@
 use PHPStan\Type\Type;
 use staabm\PHPStanDba\Error;
 
-final class LazyQueryReflector implements QueryReflector
+final class LazyQueryReflector implements QueryReflector, RecordingReflector
 {
     /**
      * @var callable():QueryReflector
@@ -51,6 +51,17 @@ public function setupDbaApi(?DbaApi $dbaApi): void
         $this->dbaApi = $dbaApi;
     }
 
+    public function getDatasource()
+    {
+        $this->reflector = $this->createReflector();
+
+        if ($this->reflector instanceof RecordingReflector) {
+            return $this->reflector->getDatasource();
+        }
+
+        return null;
+    }
+
     private function createReflector(): QueryReflector
     {
         if (null === $this->reflector) {