Improve

hakuna · Mar 26, 2024 · 8b376df · 8b376df
1 parent ce9da1c
commit 8b376df
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 52 deletions.
diff --git a/lib/rls.rb b/lib/rls.rb
@@ -26,24 +26,25 @@ def role
       configuration.role
     end
 
-    def admin=(admin)
-      RLS::Current.admin = admin
-    end
-
-    def admin
-      !!RLS::Current.admin
-    end
-
     def enable!
-      self.admin = false
+      RLS::Current.admin = false
       ActiveRecord::Base.connection_pool.disconnect!
     end
 
     def disable!
-      self.admin = true
+      RLS::Current.admin = true
       ActiveRecord::Base.connection_pool.disconnect!
     end
 
+    def without_rls(&block)
+      raise "Please supply block" unless block_given?
+
+      disable!
+      block.call
+    ensure
+      enable!
+    end
+
     def process(tenant_id, &block)
       raise "Please supply block" unless block_given?
 

diff --git a/lib/rls/extensions/postgresql_adapter.rb b/lib/rls/extensions/postgresql_adapter.rb
@@ -8,7 +8,7 @@ module PostgreSQLAdapter
 
       def initialize(...)
         super
-        execute(format(SET_ROLE_SQL, quote(RLS.role))) unless RLS.admin
+        execute(format(SET_ROLE_SQL, quote(RLS.role))) unless RLS::Current.admin
       end
 
       def rls_set(tenant_id:)

diff --git a/lib/tasks/rls.rake b/lib/tasks/rls.rake
@@ -33,44 +33,40 @@ namespace :rls do
   end
 
   task create_role: :environment do
-    RLS.disable!
-
-    RLS.connection.execute <<~SQL
-      DO $$
-      BEGIN
-        CREATE ROLE "#{RLS.role}" WITH NOLOGIN;
-      EXCEPTION
-        WHEN DUPLICATE_OBJECT THEN
-          RAISE NOTICE 'Role "#{RLS.role}" already exists';
-      END
-      $$;
-
-      GRANT ALL ON ALL TABLES IN SCHEMA public TO "#{RLS.role}";
-      GRANT ALL ON ALL SEQUENCES IN SCHEMA public TO "#{RLS.role}";
-      ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO "#{RLS.role}";
-      ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO "#{RLS.role}";
-    SQL
-
-    puts "Role #{RLS.role} created"
-
-    RLS.enable!
+    RLS.without_rls do
+      RLS.connection.execute <<~SQL
+        DO $$
+        BEGIN
+          CREATE ROLE "#{RLS.role}" WITH NOLOGIN;
+        EXCEPTION
+          WHEN DUPLICATE_OBJECT THEN
+            RAISE NOTICE 'Role "#{RLS.role}" already exists';
+        END
+        $$;
+
+        GRANT ALL ON ALL TABLES IN SCHEMA public TO "#{RLS.role}";
+        GRANT ALL ON ALL SEQUENCES IN SCHEMA public TO "#{RLS.role}";
+        ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO "#{RLS.role}";
+        ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO "#{RLS.role}";
+      SQL
+
+      puts "Role #{RLS.role} created"
+    end
   end
 
   task drop_role: :environment do
-    RLS.disable!
-
-    RLS.connection.execute <<~SQL
-      ALTER DEFAULT PRIVILEGES IN SCHEMA public REVOKE ALL ON TABLES FROM "#{RLS.role}";
-      ALTER DEFAULT PRIVILEGES IN SCHEMA public REVOKE ALL ON SEQUENCES FROM "#{RLS.role}";
-      REVOKE ALL ON ALL TABLES IN SCHEMA public FROM "#{RLS.role}";
-      REVOKE ALL ON ALL SEQUENCES IN SCHEMA public FROM "#{RLS.role}";
-      DROP OWNED BY "#{RLS.role}";
-      DROP ROLE "#{RLS.role}";
-    SQL
-
-    puts "Role #{RLS.role} dropped"
-
-    RLS.enable!
+    RLS.without_rls do
+      RLS.connection.execute <<~SQL
+        ALTER DEFAULT PRIVILEGES IN SCHEMA public REVOKE ALL ON TABLES FROM "#{RLS.role}";
+        ALTER DEFAULT PRIVILEGES IN SCHEMA public REVOKE ALL ON SEQUENCES FROM "#{RLS.role}";
+        REVOKE ALL ON ALL TABLES IN SCHEMA public FROM "#{RLS.role}";
+        REVOKE ALL ON ALL SEQUENCES IN SCHEMA public FROM "#{RLS.role}";
+        DROP OWNED BY "#{RLS.role}";
+        DROP ROLE "#{RLS.role}";
+      SQL
+
+      puts "Role #{RLS.role} dropped"
+    end
   end
 
 end
diff --git a/spec/integration/postgresql_adapter_extension_spec.rb b/spec/integration/postgresql_adapter_extension_spec.rb
@@ -12,14 +12,14 @@
 
   context "when admin" do
     specify do
-      RLS.admin = true
+      RLS::Current.admin = true
       ActiveRecord::Base.connection_pool.disconnect!
 
       role = connection.query_value("SHOW ROLE")
       expect(role).to eq "none"
 
       # ensure we roll back for remaining specs
-      RLS.admin = false
+      RLS::Current.admin = false
       ActiveRecord::Base.connection_pool.disconnect!
     end
   end

diff --git a/spec/rls_spec.rb b/spec/rls_spec.rb
@@ -14,16 +14,16 @@
   describe "#enable!" do
     subject { -> { described_class.enable! } }
 
-    before { described_class.admin = true }
+    before { RLS::Current.admin = true }
 
-    specify { expect { subject.call }.to change(described_class, :admin).from(true).to(false) }
+    specify { expect { subject.call }.to change(RLS::Current, :admin).from(true).to(false) }
     specify { expect(ActiveRecord::Base.connection_pool).to receive(:disconnect!); subject.call }
   end
 
   describe "#disable!" do
     subject { -> { described_class.disable! } }
 
-    specify { expect { subject.call }.to change(described_class, :admin).from(false).to(true) }
+    specify { expect { subject.call }.to change(RLS::Current, :admin).from(false).to(true) }
     specify { expect(ActiveRecord::Base.connection_pool).to receive(:disconnect!); subject.call }
   end
 

diff --git a/spec/unit/rls/extensions/postgresql_adapter_spec.rb b/spec/unit/rls/extensions/postgresql_adapter_spec.rb
@@ -28,7 +28,7 @@ def clear_query_cache
     end
 
     context "when admin" do
-      before { RLS.admin = true }
+      before { RLS::Current.admin = true }
 
       it "does not set the role" do
         expect_any_instance_of(MyAdapter).not_to receive(:execute).with("SET ROLE 'dummy_rls_test'")