Merge pull request #190 from mdsol/feature/accessible_objects_for_ope…

…rations_function

[MCC-833873] Add accessible_objects_for_operations function

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## 4.0.0
+* 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
+`db:migrate` to use these changes.
+
 ## 3.3.4
 * Added `fields` option to `PolicyMachineStorageAdapter::ActiveRecord` for `#accessible_objects` and
   `#accessible_objects_for_operations` to fetch only requested fields as a hash.

lib/generators/the_policy_machine/accessible_objects_for_operations_function_generator.rb

@@ -0,0 +1,12 @@
+module ThePolicyMachine
+  module Generators
+    class AccessibleObjectsForOperationsFunctionGenerator < Rails::Generators::Base
+      source_root File.expand_path('../../../migrations', __FILE__)
+
+      def generate_accessible_objects_for_operations_function_migration
+        timestamp = Time.now.utc.strftime("%Y%m%d%H%M%S")
+        copy_file('accessible_objects_for_operations_function.rb', "db/migrate/#{timestamp}_accessible_objects_for_operations_function.rb")
+      end
+    end
+  end
+end

lib/migrations/accessible_objects_for_operations_function.rb

@@ -0,0 +1,144 @@
+class AccessibleObjectsForOperationsFunction < ActiveRecord::Migration[5.2]
+  def up
+    return unless PolicyMachineStorageAdapter.postgres?
+
+    execute <<~SQL.squish
+      CREATE OR REPLACE FUNCTION pm_accessible_objects_for_operations(
+        user_id INT,
+        operations _TEXT,
+        field TEXT,
+        filters JSON DEFAULT '{}'
+      )
+      RETURNS TABLE (
+        unique_identifier varchar(255),
+        objects _varchar
+      ) AS $$
+      DECLARE
+        filter_key TEXT;
+        filter_value TEXT;
+        filter_conditions TEXT = '';
+      BEGIN
+        CREATE TEMP TABLE t_user_attribute_ids (
+          child_id INT
+        );
+
+        CREATE TEMP TABLE t_operation_set_ids (
+          operation_set_id INT,
+          object_attribute_id INT
+        );
+
+        CREATE TEMP TABLE t_accessible_operations (
+          child_id INT,
+          operation_set_id INT
+        );
+
+        CREATE TEMP TABLE t_operation_sets (
+          operation_set_id INT,
+          unique_identifier varchar(255)
+        );
+
+        SET LOCAL enable_mergejoin TO FALSE;
+
+        WITH RECURSIVE user_attribute_ids AS (
+          (
+            SELECT
+              child_id,
+              parent_id
+            FROM assignments
+            WHERE parent_id = user_id
+          )
+          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
+          )
+        )
+        INSERT INTO t_user_attribute_ids
+        SELECT child_id FROM user_attribute_ids;
+
+        IF filters IS NOT NULL AND filters::TEXT <> '{}' THEN
+          FOR filter_key, filter_value IN
+            SELECT * FROM json_each(filters)
+          LOOP
+            filter_conditions := filter_conditions || filter_key || ' = ' || filter_value || ' AND ';
+          END LOOP;
+
+          /* Chomp trailing AND */
+          filter_conditions := left(filter_conditions, -4);
+          /* Replace double quotes */
+          filter_conditions := replace(filter_conditions, '"', '''');
+
+          EXECUTE format(
+            'DELETE FROM t_user_attribute_ids ' ||
+            'WHERE child_id NOT IN (SELECT id FROM policy_elements WHERE %s)',
+            filter_conditions
+          );
+        END IF;
+
+        INSERT INTO t_operation_set_ids
+        SELECT
+          pea.operation_set_id,
+          pea.object_attribute_id
+        FROM
+          policy_element_associations pea
+          JOIN t_user_attribute_ids t ON t.child_id = pea.user_attribute_id;
+
+        WITH RECURSIVE accessible_operations AS (
+          (
+            SELECT
+              child_id,
+              parent_id AS operation_set_id
+            FROM assignments
+            WHERE parent_id IN (SELECT operation_set_id FROM t_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
+          )
+        )
+        INSERT INTO t_accessible_operations
+        SELECT * FROM accessible_operations;
+
+        INSERT INTO t_operation_sets
+        SELECT DISTINCT ao.operation_set_id, ops.unique_identifier
+          FROM
+            t_accessible_operations ao
+            JOIN policy_elements ops ON ops.id = ao.child_id
+          WHERE ops.unique_identifier = ANY (operations);
+
+        RETURN QUERY EXECUTE
+        format(
+          'SELECT os.unique_identifier, array_agg(pe.%I) AS objects ' ||
+          'FROM ' ||
+          '  t_operation_set_ids os_id ' ||
+          '  JOIN t_operation_sets os ON os.operation_set_id = os_id.operation_set_id ' ||
+          '  JOIN policy_elements pe ON pe.id = os_id.object_attribute_id ' ||
+          'WHERE pe."type" = ''PolicyMachineStorageAdapter::ActiveRecord::Object'' ' ||
+          'GROUP BY os.unique_identifier',
+          field
+        );
+
+        DROP TABLE IF EXISTS t_user_attribute_ids;
+        DROP TABLE IF EXISTS t_operation_set_ids;
+        DROP TABLE IF EXISTS t_accessible_operations;
+        DROP TABLE IF EXISTS t_operation_sets;
+
+        RETURN;
+      END;
+      $$ LANGUAGE plpgsql;
+    SQL
+  end
+
+  def down
+    execute 'DROP FUNCTION IF EXISTS pm_accessible_objects_for_operations'
+  end
+end

lib/policy_machine/version.rb

@@ -1,3 +1,3 @@
 class PolicyMachine
-  VERSION = "3.3.4"
+  VERSION = "4.0.0"
 end

lib/policy_machine_storage_adapters/active_record.rb

@@ -803,6 +803,10 @@ def batch_pluck(policy_object, query: {}, fields:, config: {}, &blk)
     # Returns all objects the user has the given operation on
     # TODO: Support multiple policy classes here
     def accessible_objects(user_or_attribute, operation, options = {})
+      if use_accessible_objects_function?(options)
+        return accessible_objects_function(user_or_attribute.id, operation, options)
+      end
+
       candidates = objects_for_user_or_attribute_and_operation(user_or_attribute, operation, options)
 
       if options[:ignore_prohibitions] || !(prohibition = prohibition_for(operation))
@@ -871,6 +875,10 @@ def accessible_objects_for_operations(user_or_attribute, operations, options = {
         raise ArgumentError, 'Functionality for indirect objects is not yet implemented!'
       end
 
+      if use_accessible_objects_function?(options)
+        return accessible_objects_for_operations_function(user_or_attribute.id, operations, options)
+      end
+
       # convert to operation names if operation instances given
       operation_names = operations.map { |o| operation_identifier(o) }
 
@@ -1077,6 +1085,25 @@ def build_accessible_object_scope(associations, options = {})
       end
     end
 
+    def use_accessible_objects_function?(options)
+      PolicyMachineStorageAdapter.postgres? &&
+        options[:direct_only] == true &&
+        options[:ignore_prohibitions] == true &&
+        options[:fields]&.one?
+    end
+
+    # Performance optimized function for PostgreSQL
+    def accessible_objects_function(user_id, operation, options)
+      accessible_objects_for_operations_function(user_id, Array.wrap(operation), options).values.flatten
+    end
+
+    # Performance optimized function for PostgreSQL
+    def accessible_objects_for_operations_function(user_id, operations, options)
+      operation_names = operations.map { |o| operation_identifier(o) }
+      accessible_map = operation_names.index_with { [] }
+      accessible_map.merge(PolicyElement.accessible_objects_for_operations(user_id, operation_names, options))
+    end
+
     def build_inclusion_scope(scope, key, value)
       Adapter.apply_include_condition(scope: scope, key: key, value: value, klass: class_for_type('object'))
     end
@@ -1317,7 +1344,7 @@ def assert_persisted_policy_element(*arguments)
     end
 
     def transaction_without_mergejoin(&block)
-      if PolicyMachineStorageAdapter::ActiveRecord::Assignment.connection.is_a? ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+      if PolicyMachineStorageAdapter.postgres?
         PolicyMachineStorageAdapter::ActiveRecord::Assignment.transaction do
           PolicyMachineStorageAdapter::ActiveRecord::Assignment.connection.execute("set local enable_mergejoin = false")
           yield
@@ -1327,4 +1354,8 @@ def transaction_without_mergejoin(&block)
       end
     end
   end
+
+  def self.postgres?
+    ::ActiveRecord::Base.connection.class.name == 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter'
+  end
 end unless active_record_unavailable

lib/policy_machine_storage_adapters/active_record/postgresql.rb

@@ -74,6 +74,37 @@ def self.operations_for_operation_sets(operation_set_ids, operation_names = nil)
         # ]
         connection.execute(sanitized_query)
       end
+
+      # The PG function can only accept a single field for now.
+      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)
+        ])
+
+        # [
+        #   { 'unique_identifier' => 'op1', 'objects' => '{obj1,obj2,obj3}' },
+        #   { 'unique_identifier' => 'op2', 'objects' => '{obj1,obj2,obj3}' },
+        # ]
+        result = connection.execute(query).to_a
+
+        # {
+        #    'op1' => ['obj1', 'obj2', 'obj3'],
+        #    'op2' => ['obj2', 'obj3', 'obj4'],
+        # }
+        decoder = PG::TextDecoder::Array.new
+        result.each_with_object({}) do |result_hash, output|
+          key = result_hash['unique_identifier']
+          objects = decoder.decode(result_hash['objects'])
+          output[key] = objects
+        end
+      end
     end
 
     class PolicyElementAssociation