Module: Apartment::Patches::ConnectionHandling

Defined in:
lib/apartment/patches/connection_handling.rb

Overview

Prepended on ActiveRecord::Base (singleton class) to intercept connection_pool lookups. When Apartment::Current.tenant is set, returns a tenant-specific pool keyed by "tenant:role", with config resolved by the adapter using the current role's base config.

Instance Method Summary collapse

Instance Method Details

#connection_pool ⇒ Object

rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity



13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
 
# File 'lib/apartment/patches/connection_handling.rb', line 13

def connection_pool # rubocop:disable Metrics/AbcSize, Metrics/MethodLength, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
  tenant = Apartment::Current.tenant
  cfg = Apartment.config

  return super if tenant.nil? || cfg.nil?
  return super if tenant.to_s == cfg.default_tenant.to_s
  return super unless Apartment.pool_manager

  # Skip tenant override for pinned models only when the adapter requires
  # a separate pool (shared_pinned_connection? is false). When shared
  # connections are supported (PG schema, MySQL), pinned models fall
  # through to the tenant pool lookup, preserving transactional integrity.
  # When adapter is nil (unconfigured), falls back to separate pool (safe default).
  adapter = Apartment.adapter
  if self != ActiveRecord::Base && Apartment.pinned_model?(self) &&
     (adapter.nil? || !adapter.shared_pinned_connection?)
    return super
  end

  # Scoped with an explicit begin rather than a method-level rescue: a
  # method-level one also covers the four `return super` guards above, so an
  # ordinary non-tenant resolution — nil tenant, most of what a Rails app does —
  # had its own failure relabelled "Failed to resolve connection pool for tenant
  # ''". A default-path resolution now raises what stock Rails raises.
  #
  # Deliberate and worth stating precisely: the boundary opens AFTER the guards,
  # so evaluating them is outside it too — `Apartment.adapter`,
  # `Apartment.pinned_model?(self)` and `adapter.shared_pinned_connection?`. A
  # raise from any of those now escapes as itself where it was previously wrapped.
  # That is the better behaviour (a broken adapter or registry is not a
  # tenant-resolution failure and should not be described as one), but it IS a
  # behaviour change, and the wrapped form is what a pre-boundary bug report would
  # have quoted.
  #
  # Inline rather than extracted into a helper for the reason the NOTE below
  # gives: the block calls `super`, which resolves only from this prepended
  # method's scope.
  begin
    # Reject pool-key-unsafe tenant names BEFORE building pool_key or entering
    # fetch_or_create. In the capped path, fetch_or_admit runs admit! (which
    # may LRU-evict an idle pool) before the adapter validates inside the
    # block, so a colon / whitespace / NUL in the raw tenant — which would
    # also corrupt the "#{tenant}:#{role}" key and PoolManager's prefix
    # matching — must be caught here. ConfigurationError is an ApartmentError,
    # so the rescue below re-raises it cleanly.
    Apartment::TenantNameValidator.validate_common!(tenant.to_s)

    role = ActiveRecord::Base.current_role
    pool_key = Apartment.pool_key(tenant, role)

    Apartment.pool_manager.fetch_or_create(pool_key) do
      # RE-ENTRANCY: when max_total_connections is set, this block runs under
      # PoolManager's @create_mutex (non-reentrant). Nothing here may resolve
      # ActiveRecord::Base.connection_pool for the current tenant — it would
      # re-enter fetch_or_create and self-deadlock. `super` resolves the
      # default pool (bypasses the patch), and check_pending_migrations? /
      # schema-cache load operate on the explicit `pool`, so all are safe.
      # Keep it that way if you add work to this block.
      # Resolve base config from the current role's default pool when available,
      # falling back to nil so the adapter uses its own base_config.
      #
      # The fallback is for a handler that has no pool for this role at all — in
      # practice a CUSTOM ConnectionHandler installed on one thread only, which is
      # this repo's own integration harness. It is NOT for parallel-migration
      # worker threads, as this comment used to claim: workers see the same
      # process-default handler and the same registrations, verified on 7.2.3.1,
      # 8.0.5 and 8.1.3.1. ddl_role is excluded from it — see
      # #guard_ddl_role_registered!.
      #
      # NOTE: `super` must be called here (not in a helper) because it refers to
      # the original connection_pool method on AR::Base, which only resolves from
      # the prepended method scope.
      base = begin
        default_pool = super
        default_pool.db_config.configuration_hash.stringify_keys
      rescue ActiveRecord::ConnectionNotEstablished
        guard_ddl_role_registered!(role, tenant)
        nil
      end

      config = Apartment.adapter.validated_connection_config(tenant, base_config_override: base)
      prefix = cfg.shard_key_prefix
      shard_key = :"#{prefix}_#{pool_key}"

      db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
        cfg.rails_env_name,
        "#{prefix}_#{pool_key}",
        config
      )

      pool = ActiveRecord::Base.connection_handler.establish_connection(
        db_config,
        owner_name: ActiveRecord::Base,
        role: role,
        shard: shard_key
      )

      # establish_connection has registered the shard in AR's ConnectionHandler.
      # If a post-establish check raises, the pool is returned to neither the
      # caller nor PoolManager — it would be orphaned: live in AR but invisible
      # to the reaper and to max_total accounting (a connection leak that also
      # undercounts the cap). Deregister it before re-raising so AR and the
      # manager stay consistent. The next request re-establishes cleanly.
      #
      # deregister_ar_shard, NOT deregister_shard: we are running inside
      # PoolManager's create block, and the full form removes from PoolManager's
      # Concurrent::Map — whose MRI backend guards compute_if_absent and delete
      # with the same non-reentrant mutex, so that would raise ThreadError
      # ("deadlock; recursive locking") and skip this deregistration entirely,
      # orphaning the very pool this rescue exists to reclaim. There is nothing
      # to remove from the manager here regardless: the pool is not stored until
      # this block returns.
      #
      # Reached with +send+ because the AR-only half is private: it is a
      # half-operation, and leaving one publicly reachable is the exact footgun
      # this seam exists to close. This is the one place it is correct.
      begin
        raise(Apartment::PendingMigrationError, tenant) if check_pending_migrations?(pool)

        load_tenant_schema_cache(tenant, pool) if cfg.schema_cache_per_tenant
      rescue StandardError
        Apartment.send(:deregister_ar_shard, pool_key)
        raise
      end

      # After the post-establish checks (a pool that fails them is discarded, so
      # extending it would be wasted), and before the pool is handed out (so the
      # very first checkin is already covered).
      # See docs/designs/transaction-taint-detection.md.
      Apartment::TransactionTaint.install(pool, tenant: tenant, pool_key: pool_key)

      pool
    end
  rescue Apartment::ApartmentError
    raise
  rescue StandardError => e
    raise(Apartment::ApartmentError,
          "Failed to resolve connection pool for tenant '#{tenant}': #{e.class}: #{e.message}")
  end
end