denormalize bearer role onto token

Author: ezekg

app/models/concerns/accountable.rb

@@ -130,8 +130,9 @@ def has_account(default: nil, **kwargs)
 
           # Perform asserts on create and update.
           validate on: %i[create update] do
+            # NB(ezekg) supports composite foreign keys, e.g. %i[bearer_type bearer_id]
             next unless
-              account_id_changed? || public_send("#{reflection.foreign_key}_changed?")
+              account_id_changed? || Array(reflection.foreign_key).any? { public_send("#{it}_changed?") }
 
             association = public_send(reflection.name)
             next if

app/models/concerns/denormalizable.rb

@@ -6,15 +6,21 @@ module Denormalizable
   DENORMALIZE_ASSOCIATION_ASYNC_BATCH_SIZE = 1_000
 
   class_methods do
-    def denormalizes(*attribute_names, with: nil, from: nil, to: nil, prefix: nil)
+    def denormalizes(*attribute_names, with: nil, from: nil, to: nil, prefix: nil, as: nil)
       raise ArgumentError, 'must provide :from, :to, or :with (but not multiple)' unless
         from.present? ^ to.present? ^ with.present?
 
+      raise ArgumentError, 'must provide either :prefix or :as (but not both)' if
+        prefix.present? && as.present?
+
+      raise ArgumentError, 'must provide a single attribute when using :as' if
+        as.present? && attribute_names.many?
+
       case
       when from.present?
-        attribute_names.each { instrument_denormalized_attribute_from(it, from:, prefix:) }
+        attribute_names.each { instrument_denormalized_attribute_from(it, from:, prefix:, as:) }
       when to.present?
-        attribute_names.each { instrument_denormalized_attribute_to(it, to:, prefix:) }
+        attribute_names.each { instrument_denormalized_attribute_to(it, to:, prefix:, as:) }
       when with.present?
         raise NotImplementedError, 'denormalizes :with is not supported yet'
       else
@@ -24,14 +30,15 @@ def denormalizes(*attribute_names, with: nil, from: nil, to: nil, prefix: nil)
 
     private
 
-    def instrument_denormalized_attribute_from(attribute_name, from:, prefix:)
+    def instrument_denormalized_attribute_from(attribute_name, from:, prefix:, as: nil)
       case from
       in Symbol => association_name if reflection = reflect_on_association(association_name)
-        prefixed_attribute_name = case prefix
-                                  when true
+        prefixed_attribute_name = case
+                                  when as.present?
+                                    as.to_s
+                                  when prefix == true
                                     "#{association_name}_#{attribute_name}"
-                                  when Symbol,
-                                       String
+                                  when (prefix in Symbol | String)
                                     "#{prefix}_#{attribute_name}"
                                   else
                                     attribute_name.to_s
@@ -41,28 +48,32 @@ def instrument_denormalized_attribute_from(attribute_name, from:, prefix:)
           raise ArgumentError, "must be a singular association: #{association_name.inspect}"
         end
 
+        # NB(ezekg) supports composite foreign keys, e.g. %i[bearer_type bearer_id]
+        association_changed = -> { Array(reflection.foreign_key).any? { send(:"#{it}_changed?") } || send(:"#{reflection.name}_changed?") }
+
         # FIXME(ezekg) after_initialize ignores prepend: false
-        set_callback :initialize, :after, -> { write_denormalized_attribute_from_schrodingers_record(association_name, attribute_name, prefixed_attribute_name) }, if: -> { send(:"#{reflection.foreign_key}_changed?") || send(:"#{reflection.name}_changed?") }, unless: :persisted?, prepend: false
-        before_validation -> { write_denormalized_attribute_from_schrodingers_record(association_name, attribute_name, prefixed_attribute_name) }, if: -> { send(:"#{reflection.foreign_key}_changed?") || send(:"#{reflection.name}_changed?") }, on: :create
-        before_update -> { write_denormalized_attribute_from_persisted_record(association_name, attribute_name, prefixed_attribute_name) }, if: -> { send(:"#{reflection.foreign_key}_changed?") || send(:"#{reflection.name}_changed?") }
+        set_callback :initialize, :after, -> { write_denormalized_attribute_from_schrodingers_record(association_name, attribute_name, prefixed_attribute_name) }, if: association_changed, unless: :persisted?, prepend: false
+        before_validation -> { write_denormalized_attribute_from_schrodingers_record(association_name, attribute_name, prefixed_attribute_name) }, if: association_changed, on: :create
+        before_update -> { write_denormalized_attribute_from_persisted_record(association_name, attribute_name, prefixed_attribute_name) }, if: association_changed
 
         # make sure validation fails if our denormalized column is modified directly
         validate -> { validate_denormalized_attribute_from_persisted_record(association_name, attribute_name, prefixed_attribute_name) }, if: :"#{prefixed_attribute_name}_changed?", on: :update
 
-        denormalized_attributes << attribute_name
+        denormalized_attributes << prefixed_attribute_name.to_sym
       else
         raise ArgumentError, "invalid :from association: #{from.inspect}"
       end
     end
 
-    def instrument_denormalized_attribute_to(attribute_name, to:, prefix:)
+    def instrument_denormalized_attribute_to(attribute_name, to:, prefix:, as: nil)
       case to
       in Symbol => association_name if reflection = reflect_on_association(association_name)
-        prefixed_attribute_name = case prefix
-                                  when true
+        prefixed_attribute_name = case
+                                  when as.present?
+                                    as.to_s
+                                  when prefix == true
                                     "#{association_name}_#{attribute_name}"
-                                  when Symbol,
-                                       String
+                                  when (prefix in Symbol | String)
                                     "#{prefix}_#{attribute_name}"
                                   else
                                     attribute_name.to_s

app/models/concerns/environmental.rb

@@ -154,8 +154,9 @@ def has_environment(default: nil, skip_verify_associations: nil, **kwargs)
 
           # Perform asserts on create and update.
           validate on: %i[create update] do
+            # NB(ezekg) supports composite foreign keys, e.g. %i[bearer_type bearer_id]
             next unless
-              environment_id_changed? || public_send("#{reflection.foreign_key}_changed?")
+              environment_id_changed? || Array(reflection.foreign_key).any? { public_send("#{it}_changed?") }
 
             association = public_send(reflection.name)
             next if

app/models/role.rb

@@ -3,6 +3,7 @@
 class Role < ApplicationRecord
   include Keygen::EE::ProtectedMethods[:permissions=, entitlements: %i[permissions]]
   include Keygen::PortableClass
+  include Denormalizable
   include Accountable
   include Dirtyable
 
@@ -36,9 +37,16 @@ class Role < ApplicationRecord
     def actions = loaded? ? collect(&:action) : super
   end
 
+  has_many :tokens,
+    primary_key: %i[resource_type resource_id],
+    foreign_key: %i[bearer_type bearer_id]
+
   # NB(ezekg) we're using account over account_id here because it may not be persisted
   has_account default: -> { resource&.account }
 
+  denormalizes :name,
+    to: :tokens, as: :bearer_role
+
   # FIXME(ezekg) replace with accountable concern i.e. an association
   delegate :default_permissions, :default_permission_ids,
     :allowed_permissions, :allowed_permission_ids,

app/models/token.rb

@@ -9,6 +9,7 @@ class Token < ApplicationRecord
   include Environmental
   include Accountable
   include Tokenable
+  include Denormalizable
   include Limitable
   include Orderable
   include Pageable
@@ -18,6 +19,10 @@ class Token < ApplicationRecord
   belongs_to :bearer,
     polymorphic: true
 
+  belongs_to :role,
+    primary_key: %i[resource_type resource_id],
+    foreign_key: %i[bearer_type bearer_id]
+
   # FIXME(ezekg) sessions must come before permissions otherwise autosave breaks
   has_many :sessions,
     dependent: :destroy_async,
@@ -46,6 +51,9 @@ def actions = loaded? ? collect(&:action) : super
     # Default to wildcard permission but allow all
     default: %w[*]
 
+  denormalizes :name,
+    from: :role, as: :bearer_role
+
   accepts_nested_attributes_for :token_permissions
   tracks_nested_attributes_for :token_permissions
 
@@ -181,15 +189,10 @@ def users        = for_bearer_type(:user)
     return none if
       names.empty?
 
-    joins(<<~SQL.squish).where(roles: { name: names })
-      INNER JOIN roles
-        ON roles.account_id    = tokens.account_id
-       AND roles.resource_type = tokens.bearer_type
-       AND roles.resource_id   = tokens.bearer_id
-    SQL
+    where(bearer_role: names)
   }
 
-  delegate :role, :role_permissions,
+  delegate :role_permissions,
     allow_nil: true,
     to: :bearer
 

db/migrate/20260707121852_add_bearer_role_to_tokens.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+class AddBearerRoleToTokens < ActiveRecord::Migration[8.1]
+  disable_ddl_transaction!
+  verbose!
+
+  def change
+    add_column :tokens, :bearer_role, :string, if_not_exists: true
+
+    add_index :tokens, %i[account_id bearer_role created_at],
+      algorithm: :concurrently,
+      if_not_exists: true
+  end
+end

db/migrate/20260707121900_seed_bearer_role_for_tokens.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+class SeedBearerRoleForTokens < ActiveRecord::Migration[8.1]
+  disable_ddl_transaction!
+  verbose!
+
+  BATCH_SIZE = 10_000
+
+  def up
+    update_count = nil
+    batch_count  = 0
+
+    until update_count == 0
+      batch_count  += 1
+      update_count  = exec_update(<<~SQL.squish, batch_count:, batch_size: BATCH_SIZE)
+        WITH batch AS (
+          SELECT
+            tokens.id  AS token_id,
+            roles.name AS role_name
+          FROM
+            tokens
+          INNER JOIN
+            roles ON roles.resource_type = tokens.bearer_type AND
+                     roles.resource_id   = tokens.bearer_id
+          WHERE
+            tokens.bearer_role IS NULL
+          LIMIT
+            :batch_size
+        )
+        UPDATE
+          tokens
+        SET
+          bearer_role = batch.role_name
+        FROM
+          batch
+        WHERE
+          tokens.id = batch.token_id
+        /* batch=:batch_count */
+      SQL
+    end
+  end
+
+  def down
+    update_count = nil
+    batch_count  = 0
+
+    until update_count == 0
+      batch_count  += 1
+      update_count  = exec_update(<<~SQL.squish, batch_count:, batch_size: BATCH_SIZE)
+        UPDATE
+          tokens
+        SET
+          bearer_role = NULL
+        WHERE
+          tokens.id IN (
+            SELECT
+              tokens.id
+            FROM
+              tokens
+            WHERE
+              tokens.bearer_role IS NOT NULL
+            LIMIT
+              :batch_size
+          )
+        /* batch=:batch_count */
+      SQL
+    end
+  end
+
+  private
+
+  def exec_update(sql, **binds)
+    ActiveRecord::Base.connection.exec_update(
+      ActiveRecord::Base.sanitize_sql([sql, **binds]),
+    )
+  end
+end

db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[8.1].define(version: 2026_05_01_131644) do
+ActiveRecord::Schema[8.1].define(version: 2026_07_07_121900) do
   # These are extensions that must be enabled in order to support this database
   enable_extension "btree_gin"
   enable_extension "pg_catalog.plpgsql"
@@ -878,6 +878,7 @@
     t.uuid "account_id"
     t.integer "activations", default: 0
     t.uuid "bearer_id"
+    t.string "bearer_role"
     t.string "bearer_type"
     t.datetime "created_at", precision: nil, null: false
     t.integer "deactivations", default: 0
@@ -888,6 +889,7 @@
     t.integer "max_deactivations"
     t.string "name"
     t.datetime "updated_at", precision: nil, null: false
+    t.index ["account_id", "bearer_role", "created_at"], name: "index_tokens_on_account_id_and_bearer_role_and_created_at"
     t.index ["account_id", "created_at"], name: "index_tokens_on_account_id_and_created_at"
     t.index ["bearer_id", "bearer_type", "created_at"], name: "index_tokens_on_bearer_id_and_bearer_type_and_created_at"
     t.index ["created_at"], name: "index_tokens_on_created_at", order: :desc

lib/tasks/keygen/permissions.rake

@@ -14,7 +14,9 @@ namespace :keygen do
       record_ids =
         args.extras.flatten.partition { Permission::ALL_PERMISSIONS.include?(it) }
 
-      records = model.includes(:account, role: { role_permissions: :permission })
+      # NB(ezekg) :role_permissions is a real association for roleable models, but
+      #           only a delegate for e.g. tokens, so this raises for non-roleables
+      records = model.includes(:account, role_permissions: :permission)
                      .where(id: record_ids)
 
       records.find_each(batch_size:) do |record|