Improve queue tenancy

src/Bootstrappers/DatabaseTenancyBootstrapper.php

@@ -6,6 +6,7 @@ namespace Stancl\Tenancy\Bootstrappers;
 
 use Stancl\Tenancy\Contracts\TenancyBootstrapper;
 use Stancl\Tenancy\Contracts\Tenant;
+use Stancl\Tenancy\Contracts\TenantWithDatabase;
 use Stancl\Tenancy\Database\DatabaseManager;
 use Stancl\Tenancy\Exceptions\TenantDatabaseDoesNotExistException;
 

src/Bootstrappers/QueueTenancyBootstrapper.php

@@ -7,7 +7,10 @@ namespace Stancl\Tenancy\Bootstrappers;
 use Illuminate\Config\Repository;
 use Illuminate\Contracts\Events\Dispatcher;
 use Illuminate\Contracts\Foundation\Application;
+use Illuminate\Queue\Events\JobFailed;
+use Illuminate\Queue\Events\JobProcessed;
 use Illuminate\Queue\Events\JobProcessing;
+use Illuminate\Queue\Events\JobRetryRequested;
 use Illuminate\Queue\QueueManager;
 use Illuminate\Support\Testing\Fakes\QueueFake;
 use Stancl\Tenancy\Contracts\TenancyBootstrapper;
@@ -28,7 +31,7 @@ class QueueTenancyBootstrapper implements TenancyBootstrapper
      */
     public static function __constructStatic(Application $app)
     {
-        static::setUpJobListener($app->make(Dispatcher::class));
+        static::setUpJobListener($app->make(Dispatcher::class), $app->runningUnitTests());
     }
 
     public function __construct(Repository $config, QueueManager $queue)
@@ -39,25 +42,70 @@ class QueueTenancyBootstrapper implements TenancyBootstrapper
         $this->setUpPayloadGenerator();
     }
 
-    protected static function setUpJobListener($dispatcher)
+    protected static function setUpJobListener($dispatcher, $runningTests)
     {
-        $dispatcher->listen(JobProcessing::class, function ($event) {
-            $tenantId = $event->job->payload()['tenant_id'] ?? null;
+        $previousTenant = null;
 
-            // The job is not tenant-aware
-            if (! $tenantId) {
-                return;
-            }
+        $dispatcher->listen(JobProcessing::class, function ($event) use (&$previousTenant) {
+            $previousTenant = tenant();
 
-            // Tenancy is already initialized for the tenant (e.g. dispatchNow was used)
-            if (tenancy()->initialized && tenant()->getTenantKey() === $tenantId) {
-                return;
-            }
-
-            // Tenancy was either not initialized, or initialized for a different tenant.
-            // Therefore, we initialize it for the correct tenant.
-            tenancy()->initialize(tenancy()->find($tenantId));
+            static::initializeTenancyForQueue($event->job->payload()['tenant_id'] ?? null);
         });
+
+        $dispatcher->listen(JobRetryRequested::class, function ($event) use (&$previousTenant) {
+            $previousTenant = tenant();
+
+            static::initializeTenancyForQueue($event->payload()['tenant_id'] ?? null);
+        });
+
+        // If we're running tests, we make sure to clean up after any artisan('queue:work') calls
+        $revertToPreviousState = function ($event) use (&$previousTenant, $runningTests) {
+            if ($runningTests) {
+                static::revertToPreviousState($event, $previousTenant);
+            }
+        };
+
+        $dispatcher->listen(JobProcessed::class, $revertToPreviousState); // artisan('queue:work') which succeeds
+        $dispatcher->listen(JobFailed::class, $revertToPreviousState); // artisan('queue:work') which fails
+    }
+
+    protected static function initializeTenancyForQueue($tenantId)
+    {
+        // The job is not tenant-aware
+        if (! $tenantId) {
+            return;
+        }
+
+        if (tenancy()->initialized) {
+            if (tenant()->getTenantKey() === $tenantId) {
+                // Tenancy is already initialized for the tenant (e.g. dispatchNow was used)
+                return;
+            }
+        }
+
+        // Tenancy was either not initialized, or initialized for a different tenant.
+        // Therefore, we initialize it for the correct tenant.
+        tenancy()->initialize(tenancy()->find($tenantId));
+    }
+
+    protected static function revertToPreviousState($event, &$previousTenant)
+    {
+        $tenantId = $event->job->payload()['tenant_id'] ?? null;
+
+        // The job was not tenant-aware
+        if (! $tenantId) {
+            return;
+        }
+
+        // Revert back to the previous tenant
+        if (tenant() && $previousTenant && $previousTenant->isNot(tenant())) {
+            tenancy()->initialize($previousTenant);
+        }
+
+        // End tenancy
+        if (tenant() && (! $previousTenant)) {
+            tenancy()->end();
+        }
     }
 
     protected function setUpPayloadGenerator()

src/Database/DatabaseManager.php

@@ -38,7 +38,7 @@ class DatabaseManager
      */
     public function connectToTenant(TenantWithDatabase $tenant)
     {
-        $this->database->purge('tenant');
+        $this->purgeTenantConnection();
         $this->createTenantConnection($tenant);
         $this->setDefaultConnection('tenant');
     }
@@ -48,10 +48,7 @@ class DatabaseManager
      */
     public function reconnectToCentral()
     {
-        if (tenancy()->initialized) {
-            $this->database->purge('tenant');
-        }
-
+        $this->purgeTenantConnection();
         $this->setDefaultConnection($this->config->get('tenancy.database.central_connection'));
     }
 
@@ -60,7 +57,7 @@ class DatabaseManager
      */
     public function setDefaultConnection(string $connection)
     {
-        $this->app['config']['database.default'] = $connection;
+        $this->config['database.default'] = $connection;
         $this->database->setDefaultConnection($connection);
     }
 
@@ -69,7 +66,19 @@ class DatabaseManager
      */
     public function createTenantConnection(TenantWithDatabase $tenant)
     {
-        $this->app['config']['database.connections.tenant'] = $tenant->database()->connection();
+        $this->config['database.connections.tenant'] = $tenant->database()->connection();
+    }
+
+    /**
+     * Purge the tenant database connection.
+     */
+    public function purgeTenantConnection()
+    {
+        if (array_key_exists('tenant', $this->database->getConnections())) {
+            $this->database->purge('tenant');
+        }
+
+        unset($this->config['database.connections.tenant']);
     }
 
     /**

src/Exceptions/TenantCouldNotBeIdentifiedById.php

@@ -9,7 +9,7 @@ use Facade\IgnitionContracts\ProvidesSolution;
 use Facade\IgnitionContracts\Solution;
 use Stancl\Tenancy\Contracts\TenantCouldNotBeIdentifiedException;
 
-// todo: in v3 this should be suffixed with Exception
+// todo: in v4 this should be suffixed with Exception
 class TenantCouldNotBeIdentifiedById extends TenantCouldNotBeIdentifiedException implements ProvidesSolution
 {
     public function __construct($tenant_id)

src/Tenancy.php

@@ -66,10 +66,10 @@ class Tenancy
             return;
         }
 
-        $this->initialized = false;
-
         event(new Events\TenancyEnded($this));
 
+        $this->initialized = false;
+
         $this->tenant = null;
     }