Improve queue tenancy

docker-compose.yml

@@ -4,7 +4,7 @@ services:
         build:
             context: .
             args:
-              PHP_VERSION: ${PHP_VERSION}
+              PHP_VERSION: ${PHP_VERSION:-8.1}
         depends_on:
             mysql:
                 condition: service_healthy

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;
     }
 

tests/QueueTest.php

@@ -4,18 +4,30 @@ declare(strict_types=1);
 
 namespace Stancl\Tenancy\Tests;
 
+use Exception;
 use Illuminate\Bus\Queueable;
 use Illuminate\Contracts\Queue\ShouldQueue;
+use Illuminate\Database\Schema\Blueprint;
 use Illuminate\Foundation\Bus\Dispatchable;
+use Illuminate\Queue\Events\JobProcessed;
 use Illuminate\Queue\Events\JobProcessing;
 use Illuminate\Queue\InteractsWithQueue;
 use Illuminate\Queue\SerializesModels;
+use Illuminate\Support\Facades\DB;
 use Illuminate\Support\Facades\Event;
+use Illuminate\Support\Facades\Schema;
 use Spatie\Valuestore\Valuestore;
+use Stancl\JobPipeline\JobPipeline;
+use Stancl\Tenancy\Bootstrappers\DatabaseTenancyBootstrapper;
 use Stancl\Tenancy\Bootstrappers\QueueTenancyBootstrapper;
+use Stancl\Tenancy\Events\TenancyEnded;
 use Stancl\Tenancy\Events\TenancyInitialized;
+use Stancl\Tenancy\Events\TenantCreated;
+use Stancl\Tenancy\Jobs\CreateDatabase;
 use Stancl\Tenancy\Listeners\BootstrapTenancy;
+use Stancl\Tenancy\Listeners\RevertToCentralContext;
 use Stancl\Tenancy\Tests\Etc\Tenant;
+use Stancl\Tenancy\Tests\Etc\User;
 
 class QueueTest extends TestCase
 {
@@ -31,15 +43,49 @@ class QueueTest extends TestCase
         config([
             'tenancy.bootstrappers' => [
                 QueueTenancyBootstrapper::class,
+                DatabaseTenancyBootstrapper::class,
             ],
             'queue.default' => 'redis',
         ]);
 
         Event::listen(TenancyInitialized::class, BootstrapTenancy::class);
+        Event::listen(TenancyEnded::class, RevertToCentralContext::class);
 
         $this->valuestore = Valuestore::make(__DIR__ . '/Etc/tmp/queuetest.json')->flush();
     }
 
+    protected function withFailedJobs()
+    {
+        Schema::connection('central')->create('failed_jobs', function (Blueprint $table) {
+            $table->id();
+            $table->string('uuid')->unique();
+            $table->text('connection');
+            $table->text('queue');
+            $table->longText('payload');
+            $table->longText('exception');
+            $table->timestamp('failed_at')->useCurrent();
+        });
+    }
+
+    protected function withUsers()
+    {
+        Schema::create('users', function (Blueprint $table) {
+            $table->increments('id');
+            $table->string('name');
+            $table->string('email')->unique();
+            $table->string('password');
+            $table->rememberToken();
+            $table->timestamps();
+        });
+    }
+
+    protected function withTenantDatabases()
+    {
+        Event::listen(TenantCreated::class, JobPipeline::make([CreateDatabase::class])->send(function (TenantCreated $event) {
+            return $event->tenant;
+        })->toListener());
+    }
+
     /** @test */
     public function tenant_id_is_passed_to_tenant_queues()
     {
@@ -49,7 +95,7 @@ class QueueTest extends TestCase
 
         tenancy()->initialize($tenant);
 
-        Event::fake([JobProcessing::class]);
+        Event::fake([JobProcessing::class, JobProcessed::class]);
 
         dispatch(new TestJob($this->valuestore));
 
@@ -79,21 +125,91 @@ class QueueTest extends TestCase
         });
     }
 
-    /** @test */
-    public function tenancy_is_initialized_inside_queues()
+    /**
+     * @test
+     *
+     * @testWith [true]
+     *           [false]
+     */
+    public function tenancy_is_initialized_inside_queues(bool $shouldEndTenancy)
     {
-        $tenant = Tenant::create([
-            'id' => 'acme',
-        ]);
+        $this->withTenantDatabases();
+        $this->withFailedJobs();
+
+        $tenant = Tenant::create();
 
         tenancy()->initialize($tenant);
 
-        dispatch(new TestJob($this->valuestore));
+        $this->withUsers();
+
+        $user = User::create(['name' => 'Foo', 'email' => 'foo@bar.com', 'password' => 'secret']);
+
+        $this->valuestore->put('userName', 'Bar');
+
+        dispatch(new TestJob($this->valuestore, $user));
 
         $this->assertFalse($this->valuestore->has('tenant_id'));
+
+        if ($shouldEndTenancy) {
+            tenancy()->end();
+        }
+
         $this->artisan('queue:work --once');
 
-        $this->assertSame('The current tenant id is: acme', $this->valuestore->get('tenant_id'));
+        $this->assertSame(0, DB::connection('central')->table('failed_jobs')->count());
+
+        $this->assertSame('The current tenant id is: ' . $tenant->id, $this->valuestore->get('tenant_id'));
+
+        $tenant->run(function () use ($user) {
+            $this->assertSame('Bar', $user->fresh()->name);
+        });
+    }
+
+    /**
+     * @test
+     *
+     * @testWith [true]
+     *           [false]
+     */
+    public function tenancy_is_initialized_when_retrying_jobs(bool $shouldEndTenancy)
+    {
+        $this->withFailedJobs();
+        $this->withTenantDatabases();
+
+        $tenant = Tenant::create();
+
+        tenancy()->initialize($tenant);
+
+        $this->withUsers();
+
+        $user = User::create(['name' => 'Foo', 'email' => 'foo@bar.com', 'password' => 'secret']);
+
+        $this->valuestore->put('userName', 'Bar');
+        $this->valuestore->put('shouldFail', true);
+
+        dispatch(new TestJob($this->valuestore, $user));
+
+        $this->assertFalse($this->valuestore->has('tenant_id'));
+
+        if ($shouldEndTenancy) {
+            tenancy()->end();
+        }
+
+        $this->artisan('queue:work --once');
+
+        $this->assertSame(1, DB::connection('central')->table('failed_jobs')->count());
+        $this->assertNull($this->valuestore->get('tenant_id')); // job failed
+
+        $this->artisan('queue:retry all');
+        $this->artisan('queue:work --once');
+
+        $this->assertSame(0, DB::connection('central')->table('failed_jobs')->count());
+
+        $this->assertSame('The current tenant id is: ' . $tenant->id, $this->valuestore->get('tenant_id')); // job succeeded
+
+        $tenant->run(function () use ($user) {
+            $this->assertSame('Bar', $user->fresh()->name);
+        });
     }
 
     /** @test */
@@ -127,13 +243,31 @@ class TestJob implements ShouldQueue
     /** @var Valuestore */
     protected $valuestore;
 
-    public function __construct(Valuestore $valuestore)
+    /** @var User|null */
+    protected $user;
+
+    public function __construct(Valuestore $valuestore, User $user = null)
     {
         $this->valuestore = $valuestore;
+        $this->user = $user;
     }
 
     public function handle()
     {
+        if ($this->valuestore->get('shouldFail')) {
+            $this->valuestore->put('shouldFail', false);
+
+            throw new Exception('failing');
+        }
+
+        if ($this->user) {
+            assert($this->user->getConnectionName() === 'tenant');
+        }
+
         $this->valuestore->put('tenant_id', 'The current tenant id is: ' . tenant('id'));
+
+        if ($userName = $this->valuestore->get('userName')) {
+            $this->user->update(['name' => $userName]);
+        }
     }
 }

tests/TenantDatabaseManagerTest.php

@@ -4,17 +4,21 @@ declare(strict_types=1);
 
 namespace Stancl\Tenancy\Tests;
 
+use Illuminate\Database\Schema\Blueprint;
 use Illuminate\Support\Facades\Event;
+use Illuminate\Support\Facades\Schema;
 use Illuminate\Support\Str;
 use PDO;
 use Stancl\JobPipeline\JobPipeline;
 use Stancl\Tenancy\Bootstrappers\DatabaseTenancyBootstrapper;
 use Stancl\Tenancy\Database\DatabaseManager;
+use Stancl\Tenancy\Events\TenancyEnded;
 use Stancl\Tenancy\Events\TenancyInitialized;
 use Stancl\Tenancy\Events\TenantCreated;
 use Stancl\Tenancy\Exceptions\TenantDatabaseAlreadyExistsException;
 use Stancl\Tenancy\Jobs\CreateDatabase;
 use Stancl\Tenancy\Listeners\BootstrapTenancy;
+use Stancl\Tenancy\Listeners\RevertToCentralContext;
 use Stancl\Tenancy\TenantDatabaseManagers\MySQLDatabaseManager;
 use Stancl\Tenancy\TenantDatabaseManagers\PermissionControlledMySQLDatabaseManager;
 use Stancl\Tenancy\TenantDatabaseManagers\PostgreSQLDatabaseManager;
@@ -102,6 +106,52 @@ class TenantDatabaseManagerTest extends TestCase
         ];
     }
 
+    /** @test */
+    public function the_tenant_connection_is_fully_removed()
+    {
+        config([
+            'tenancy.boostrappers' => [
+                DatabaseTenancyBootstrapper::class,
+            ],
+        ]);
+
+        Event::listen(TenantCreated::class, JobPipeline::make([CreateDatabase::class])->send(function (TenantCreated $event) {
+            return $event->tenant;
+        })->toListener());
+
+        Event::listen(TenancyInitialized::class, BootstrapTenancy::class);
+        Event::listen(TenancyEnded::class, RevertToCentralContext::class);
+
+        $tenant = Tenant::create();
+
+        $this->assertSame(['central'], array_keys(app('db')->getConnections()));
+        $this->assertArrayNotHasKey('tenant', config('database.connections'));
+
+        tenancy()->initialize($tenant);
+
+        $this->createUsersTable();
+
+        $this->assertSame(['central', 'tenant'], array_keys(app('db')->getConnections()));
+        $this->assertArrayHasKey('tenant', config('database.connections'));
+
+        tenancy()->end();
+
+        $this->assertSame(['central'], array_keys(app('db')->getConnections()));
+        $this->assertNull(config('database.connections.tenant'));
+    }
+
+    protected function createUsersTable()
+    {
+        Schema::create('users', function (Blueprint $table) {
+            $table->increments('id');
+            $table->string('name');
+            $table->string('email')->unique();
+            $table->string('password');
+            $table->rememberToken();
+            $table->timestamps();
+        });
+    }
+
     /** @test */
     public function db_name_is_prefixed_with_db_path_when_sqlite_is_used()
     {
@@ -217,5 +267,6 @@ class TenantDatabaseManagerTest extends TestCase
     /** @test */
     public function path_used_by_sqlite_manager_can_be_customized()
     {
+        $this->markTestIncomplete();
     }
 }