[MCC-833873] Fix single pg call for accessible_objects from replica

CHANGELOG.md

@@ -1,10 +1,13 @@
 # Changelog
 
+## 4.0.2
+* Updated `accessible_objects_for_operations` code path to use a CTE when using a replica database.
+
 ## 4.0.1
-* Updated PostgreSQL function `pm_accessible_objects_for_operations` to ensure uniqueness of results
+* Updated PostgreSQL function `pm_accessible_objects_for_operations` to ensure uniqueness of results.
 
 ## 4.0.0
-* Added a PostgreSQL function for `#accessible_objects` and `#accessible_objects_for_operations` which are performance
+* Added a PostgreSQL function for `#accessible_objects` and `#accessible_objects_for_operations` which are performance.
 optimized. Only supported for a single `field`, `direct_only`, and `ignore_prohibitions`.
 
 – Execute `bundle exec rails generate the_policy_machine:accessible_objects_for_operations_function` and rerun

lib/policy_machine/version.rb

@@ -1,3 +1,3 @@
 class PolicyMachine
-  VERSION = "4.0.1"
+  VERSION = "4.0.2"
 end

lib/policy_machine_storage_adapters/active_record/postgresql.rb

@@ -80,13 +80,22 @@ def self.accessible_objects_for_operations(user_id, operation_names, options)
         field = options[:fields].first
         filters = options.dig(:filters, :user_attributes) || {}
 
-        query = sanitize_sql_for_assignment([
-          'SELECT * FROM pm_accessible_objects_for_operations(?,?,?,?)',
-          user_id,
-          PG::TextEncoder::Array.new.encode(operation_names),
-          field,
-          JSON.dump(filters)
-        ])
+        query =
+          if replica?
+            sanitize_sql_for_assignment([
+              accessible_objects_for_operations_cte(field, filters),
+              user_id,
+              operation_names
+            ])
+          else
+            sanitize_sql_for_assignment([
+              'SELECT * FROM pm_accessible_objects_for_operations(?,?,?,?)',
+              user_id,
+              PG::TextEncoder::Array.new.encode(operation_names),
+              field,
+              JSON.dump(filters)
+            ])
+          end
 
         # [
         #   { 'unique_identifier' => 'op1', 'objects' => '{obj1,obj2,obj3}' },
@@ -105,6 +114,112 @@ def self.accessible_objects_for_operations(user_id, operation_names, options)
           output[key] = objects
         end
       end
+
+      private
+
+      # For replica database connections which do not support temporary tables.
+      # This is a little slower than the PG function but still quicker than the
+      # existing ActiveRecord code path.
+      def self.accessible_objects_for_operations_cte(field, filters)
+        <<~SQL.squish
+          SET LOCAL enable_mergejoin TO FALSE;
+
+          WITH RECURSIVE "user" AS (
+            SELECT *
+            FROM policy_elements
+            WHERE
+              id = ?
+          ),
+          user_attribute_ids AS (
+            (
+              SELECT
+                child_id,
+                parent_id
+              FROM assignments
+              WHERE parent_id = (SELECT id FROM "user")
+            )
+            UNION ALL
+            (
+              SELECT
+                a.child_id,
+                a.parent_id
+              FROM
+                assignments a
+                JOIN user_attribute_ids ua_id ON ua_id.child_id = a.parent_id
+            )
+          ),
+          operation_set_ids AS (
+            SELECT
+              operation_set_id,
+              object_attribute_id
+            FROM policy_element_associations
+            WHERE user_attribute_id IN (SELECT child_id FROM user_attribute_ids #{accessible_object_filters(filters)})
+          ),
+          accessible_operations AS (
+            (
+              SELECT
+                child_id,
+                parent_id AS operation_set_id
+              FROM assignments
+              WHERE parent_id IN (SELECT operation_set_id FROM operation_set_ids)
+            )
+            UNION ALL
+            (
+              SELECT
+                a.child_id,
+                op.operation_set_id AS operation_set_id
+              FROM
+                assignments a
+                JOIN accessible_operations op ON op.child_id = a.parent_id
+            )
+          ),
+          operation_sets AS (
+            SELECT DISTINCT ao.operation_set_id, ops.unique_identifier
+            FROM
+              accessible_operations ao
+              JOIN policy_elements ops ON ops.id = ao.child_id
+            WHERE
+              ops.unique_identifier IN (?)
+          ),
+          operation_objects AS (
+            SELECT
+              os.unique_identifier,
+              array_remove(array_agg(
+                (
+                  SELECT pe.#{connection.quote_column_name(field)}
+                  FROM policy_elements pe
+                  WHERE
+                    pe.id = os_id.object_attribute_id
+                    AND "type" = 'PolicyMachineStorageAdapter::ActiveRecord::Object'
+                )
+              ), NULL) AS objects
+            FROM
+              operation_set_ids os_id
+              JOIN operation_sets os ON os.operation_set_id = os_id.operation_set_id
+            GROUP BY os.unique_identifier
+          )
+          SELECT
+            unique_identifier,
+            ARRAY(SELECT DISTINCT o FROM UNNEST(objects) AS a(o)) as objects
+          FROM operation_objects;
+        SQL
+      end
+
+      def self.accessible_object_filters(filters)
+        return '' if filters.blank?
+
+        condition = 'WHERE child_id IN (SELECT id FROM policy_elements WHERE '
+
+        filters.each do |key, value|
+          condition << sanitize_sql_for_assignment(["#{connection.quote_column_name(key)} = ? AND ", value])
+        end
+
+        condition.chomp('AND ') << ')'
+      end
+
+      def self.replica?
+        ::ActiveRecord::Base.connection_db_config.configuration_hash[:replica] == true
+      end
     end
 
     class PolicyElementAssociation

policy_machine.gemspec

@@ -33,5 +33,5 @@ Gem::Specification.new do |s|
   s.add_development_dependency('byebug')
   s.add_development_dependency('neography', '~> 1.1')
   s.add_development_dependency('database_cleaner')
-  s.add_development_dependency('rails', '~> 6.0')
+  s.add_development_dependency('rails', '~> 6.1')
 end

spec/policy_machine_storage_adapters/active_record_spec.rb

@@ -741,25 +741,12 @@
         end
       end
 
-      describe 'accessible_objects_for_operations_function' do
+      shared_examples 'a single query for accessible objects' do
         before do
           priv_pm.add_association(color_1, creator, object_6)
           priv_pm.add_association(color_2, creator, object_7)
         end
 
-        it 'uses a PostgreSQL function' do
-          expect_any_instance_of(PolicyMachineStorageAdapter::ActiveRecord)
-            .to receive(:accessible_objects_for_operations_function)
-
-          priv_pm.accessible_objects_for_operations(
-            user_1,
-            [create, paint],
-            direct_only: true,
-            ignore_prohibitions: true,
-            fields: [:unique_identifier]
-          )
-        end
-
         it 'returns objects accessible via each of multiple given operations' do
           result = priv_pm.accessible_objects_for_operations(
             user_1,
@@ -859,6 +846,46 @@
         end
       end
 
+      describe 'accessible_objects_for_operations_function' do
+        it_behaves_like 'a single query for accessible objects'
+
+        it 'uses a PostgreSQL function' do
+          expect_any_instance_of(PolicyMachineStorageAdapter::ActiveRecord)
+            .to receive(:accessible_objects_for_operations_function)
+
+          priv_pm.accessible_objects_for_operations(
+            user_1,
+            [create, paint],
+            direct_only: true,
+            ignore_prohibitions: true,
+            fields: [:unique_identifier]
+          )
+        end
+      end
+
+      describe 'accessible_objects_for_operations_cte' do
+        before do
+          allow(PolicyMachineStorageAdapter::ActiveRecord::PolicyElement)
+            .to receive(:replica?).and_return(true)
+        end
+
+        it_behaves_like 'a single query for accessible objects'
+
+        it 'uses a single CTE' do
+          expect(PolicyMachineStorageAdapter::ActiveRecord::PolicyElement)
+            .to receive(:accessible_objects_for_operations_cte)
+            .and_call_original
+
+          priv_pm.accessible_objects_for_operations(
+            user_1,
+            [create, paint],
+            direct_only: true,
+            ignore_prohibitions: true,
+            fields: [:unique_identifier]
+          )
+        end
+      end
+
       describe 'accessible_ancestor_objects' do
         context 'when policy element associations are not provided as an argument' do
           it 'returns objects accessible via the filtered attribute on an object scope' do