Feat/handling retries on sidekiq & set admin user before migrations and rollbacks

lib/pg_rls/database/tasks/admin_database.rake

@@ -22,14 +22,13 @@ end
 namespace :db do
   include PgRls::Schema::UpStatements
 
-  override_task grant_usage: :load_config do
+  override_task :load_config do
     PgRls.instance_variable_set(:@as_db_admin, true)
-    create_rls_user
+    Rake::Task['db:load_config:original'].invoke
   end
 
-  override_task abort_if_pending_migrations: :load_config do
-    PgRls.instance_variable_set(:@as_db_admin, true)
-    Rake::Task['db:abort_if_pending_migrations:original'].invoke
+  override_task grant_usage: :load_config do
+    create_rls_user
   end
 
   namespace :test do

lib/pg_rls/middleware/sidekiq/client.rb

@@ -12,12 +12,9 @@ def call(_job_class, msg, _queue, _redis_pool)
         end
 
         def load_tenant_attribute!(msg)
-          if PgRls.admin_connection?
-            msg['admin'] = true
-          else
-            tenant = PgRls::Tenant.fetch!
-            msg['pg_rls'] = tenant.id
-          end
+          return msg['admin'] = true if PgRls.admin_connection?
+
+          msg['pg_rls'] ||= PgRls::Tenant.fetch&.id
         end
       end
     end

spec/middleware/sidekiq/client_spec.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+require 'pg_rls/middleware/sidekiq/client'
+require 'sidekiq/testing'
+
+# Define the dummy workers
+class DummyAdminWorker
+  include Sidekiq::Worker
+
+  def perform(*_args)
+    raise 'Not an admin connection' unless PgRls.admin_connection?
+
+    'Admin task'
+  end
+end
+
+class DummyTenantWorker
+  include Sidekiq::Worker
+
+  def perform(*_args)
+    raise 'Not a tenant connection' if PgRls.admin_connection?
+
+    PgRls::Tenant.fetch
+  end
+end
+
+class DummyFailingWorker
+  include Sidekiq::Worker
+
+  def perform(*_args)
+    raise 'Failed job'
+  end
+end
+
+RSpec.describe PgRls::Middleware::Sidekiq::Client do
+  before(:all) do
+    Sidekiq.configure_client do |config|
+      config.logger.level = Logger::WARN
+      config.client_middleware do |chain|
+        chain.add described_class
+      end
+    end
+  end
+
+  describe 'middleware behavior' do
+    context 'when admin connection is true' do
+      before do
+        allow(PgRls).to receive(:admin_connection?).and_return(true)
+      end
+
+      it 'sets admin attribute to true for the job' do
+        Sidekiq::Testing.inline! do
+          DummyAdminWorker.perform_async
+        end
+      end
+    end
+
+    context 'when admin connection is false' do
+      let(:tenant) { double('Tenant', id: 123) }
+
+      before do
+        allow(PgRls).to receive(:admin_connection?).and_return(false)
+        allow(PgRls::Tenant).to receive(:fetch).and_return(tenant)
+      end
+
+      it 'sets pg_rls attribute with tenant id for the job' do
+        Sidekiq::Testing.inline! do
+          DummyTenantWorker.perform_async
+        end
+      end
+    end
+
+    context 'when the job fails' do
+      let(:tenant) { double('Tenant', id: 123) }
+
+      before do
+        allow(PgRls).to receive(:admin_connection?).and_return(false)
+        allow(PgRls::Tenant).to receive(:fetch).and_return(tenant)
+      end
+
+      it 'raises the error triggered by the job (not other raised in the middleware)' do
+        Sidekiq::Testing.inline! do
+          expect { DummyFailingWorker.perform_async }.to raise_error('Failed job')
+        end
+      end
+    end
+  end
+end

spec/pg_rls_spec.rb

@@ -6,6 +6,6 @@
   end
 
   it 'does something useful' do
-    expect(require_relative('lib/PgRls/version')).to be(true)
+    expect(require_relative('../lib/pg_rls/version')).to be(false)
   end
 end