Class: PostgresMultitenant::SchemaManager

Inherits:
Object
  • Object
show all
Defined in:
lib/postgres_multitenant/schema_manager.rb

Constant Summary collapse

RESERVED_SCHEMAS = 
%w[public information_schema].freeze
SYSTEM_SCHEMA_PREFIX = 
'pg_'

Class Method Summary collapse

Class Method Details

.cleanup_orphaned_schemas(dry_run: false) ⇒ Hash

Clean up all orphaned schemas (schemas without corresponding organizations)

Parameters:

  • dry_run (Boolean) (defaults to: false) —

    If true, only report what would be cleaned (default: false)

Returns:

  • (Hash) —

    Cleanup report with :cleaned, :failed, or :would_clean keys



267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 267

def cleanup_orphaned_schemas(dry_run: false)
  orphaned = find_orphaned_schemas

  if orphaned.empty?
    Rails.logger.info('No orphaned schemas found')
    return { cleaned: [], failed: [], would_clean: [] }
  end

  Rails.logger.info("Found #{orphaned.size} orphaned schema#{'s' if orphaned.size != 1}: #{orphaned.join(', ')}")

  if dry_run
    Rails.logger.info("DRY RUN: Would clean up #{orphaned.size} schema#{'s' if orphaned.size != 1}")
    return { would_clean: orphaned }
  end

  cleaned = []
  failed = []

  orphaned.each do |schema_name|
    begin
      drop_schema_with_retry(schema_name, raise_on_failure: true)
      Rails.logger.info("Successfully cleaned up orphaned schema: #{schema_name}")
      cleaned << schema_name
    rescue StandardError => e
      Rails.logger.error("Failed to clean up orphaned schema #{schema_name}: #{e.message}")
      failed << schema_name
    end
  end

  Rails.logger.info("Cleanup complete. Cleaned: #{cleaned.size}, Failed: #{failed.size}")
  { cleaned: cleaned, failed: failed }
end

.column_exists_in_schema?(table_name, column_name) ⇒ Boolean

Helper method to check if a column exists in a table

Returns:

  • (Boolean) 


58
59
60
61
62
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 58

def self.column_exists_in_schema?(table_name, column_name)
  connection.column_exists?(table_name, column_name)
rescue ActiveRecord::StatementInvalid
  false
end

.create_schema(name) ⇒ Object

Create a new tenant schema and run migrations



70
71
72
73
74
75
76
77
78
79
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 70

def create_schema(name)
  validate_schema_name!(name)

  connection.execute("CREATE SCHEMA IF NOT EXISTS #{connection.quote_table_name(name)}")

  # Load schema.rb into the new tenant schema
  with_schema(name) do
    load_schema_file
  end
end

.current_schema ⇒ Object

Get current schema from search_path



104
105
106
107
108
109
110
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 104

def current_schema
  conn = connection
  ensure_schema_search_path_supported!(conn)

  schema = conn.schema_search_path.to_s.split(',').first.to_s.strip.delete('"')
  schema.presence || 'public'
end

.drop_schema(name) ⇒ Object

Drop a tenant schema



82
83
84
85
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 82

def drop_schema(name)
  validate_schema_name!(name)
  connection.execute("DROP SCHEMA IF EXISTS #{connection.quote_table_name(name)} CASCADE")
end

.drop_schema_with_retry(name, max_attempts: 5, raise_on_failure: true) ⇒ Boolean

Drop a schema with retry logic and connection termination

Parameters:

  • name (String) —

    The schema name to drop

  • max_attempts (Integer) (defaults to: 5) —

    Maximum number of retry attempts (default: 5)

  • raise_on_failure (Boolean) (defaults to: true) —

    Whether to raise exception on failure (default: true)

Returns:

  • (Boolean) —

    true if successful, false if failed and raise_on_failure is false

Raises:

  • (SchemaDropError) —

    if max attempts exceeded and raise_on_failure is true



146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 146

def drop_schema_with_retry(name, max_attempts: 5, raise_on_failure: true)
  validate_schema_name!(name)

  # If schema doesn't exist, return true (idempotent)
  unless schema_exists?(name)
    Rails.logger.info("Schema #{name} does not exist, skipping drop")
    return true
  end

  attempts = 0
  last_error = nil

  while attempts < max_attempts
    attempts += 1

    begin
      # Terminate active connections only when explicitly configured.
      terminate_schema_connections(name) if PostgresMultitenant.configuration.terminate_connections_on_drop

      Rails.logger.info("Dropping schema: #{name} (attempt #{attempts}/#{max_attempts})")
      drop_schema(name)
      Rails.logger.info("Successfully dropped schema: #{name}")
      return true
    rescue ActiveRecord::StatementInvalid => e
      last_error = e

      # Check if error is retryable (locked schema, active connections, etc.)
      if e.message.include?('PG::ObjectInUse') ||
         e.message.include?('being accessed by other users') ||
         e.message.include?('deadlock') ||
         e.message.include?('lock timeout')

        if attempts < max_attempts
          wait_time = 2**(attempts - 1) # Exponential backoff: 1, 2, 4, 8, 16
          Rails.logger.warn(
            "Failed to drop schema #{name} (attempt #{attempts}/#{max_attempts}): " \
            "#{e.message}. Retrying in #{wait_time}s..."
          )
          sleep(wait_time)
        else
          # Max attempts reached
          Rails.logger.error("Failed to drop schema #{name} after #{max_attempts} attempts: #{e.message}")
          if raise_on_failure
            raise SchemaDropError,
                  "Failed to drop schema '#{name}' after #{max_attempts} attempts: #{e.message}"
          else
            return false
          end
        end
      else
        # Non-retryable error
        Rails.logger.error("Non-retryable error dropping schema #{name}: #{e.message}")
        if raise_on_failure
          raise SchemaDropError, "Failed to drop schema '#{name}': #{e.message}"
        else
          return false
        end
      end
    end
  end

  # Should not reach here, but handle it
  if raise_on_failure
    raise SchemaDropError, "Failed to drop schema '#{name}' after #{max_attempts} attempts: #{last_error&.message}"
  else
    false
  end
end

.find_orphaned_schemas ⇒ Array<String>

Find schemas that don't have a corresponding organization

Returns:

  • (Array<String>) —

    List of orphaned schema names



241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 241

def find_orphaned_schemas
  all_tenant_schemas = tenant_schemas
  tenant_class = PostgresMultitenant.tenant_class
  return [] unless tenant_class

  tenant_schema_names = []
  if tenant_class.respond_to?(:find_each)
    tenant_class.find_each do |tenant|
      schema_name = PostgresMultitenant.schema_name_for(tenant)
      tenant_schema_names << schema_name if schema_name.present?
    end
  else
    tenant_class.all.each do |tenant|
      schema_name = PostgresMultitenant.schema_name_for(tenant)
      tenant_schema_names << schema_name if schema_name.present?
    end
  end

  orphaned = all_tenant_schemas - tenant_schema_names
  orphaned
end

.reset! ⇒ Object

Reset connection to public schema



135
136
137
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 135

def reset!
  switch_to('public')
end

.reset_tenant_table_cache! ⇒ Object 



64
65
66
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 64

def self.reset_tenant_table_cache!
  @tenant_tables_cache = nil
end

.schema_exists?(name) ⇒ Boolean

Check if a schema exists

Returns:

  • (Boolean) 


124
125
126
127
128
129
130
131
132
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 124

def schema_exists?(name)
  result = connection.select_value(
    ActiveRecord::Base.sanitize_sql_array([
                                            'SELECT 1 FROM information_schema.schemata WHERE schema_name = ? LIMIT 1',
                                            name
                                          ])
  )
  result.present?
end

.tenant_schemas ⇒ Object

List all tenant schemas (excluding system schemas)



113
114
115
116
117
118
119
120
121
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 113

def tenant_schemas
  connection.select_values(
    ActiveRecord::Base.sanitize_sql_array([
                                            'SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN (?) AND schema_name NOT LIKE ?',
                                            RESERVED_SCHEMAS,
                                            "#{SYSTEM_SCHEMA_PREFIX}%"
                                          ])
  )
end

.tenant_tables ⇒ Object

SECURITY: Only tenant-specific tables should be created in tenant schemas.

Raises:

  • (StandardError) 


16
17
18
19
20
21
22
23
24
25
26
27
28
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 16

def self.tenant_tables
  return @tenant_tables_cache if defined?(@tenant_tables_cache) && @tenant_tables_cache

  resolver = PostgresMultitenant.configuration.tenant_table_resolver
  raise StandardError, 'Configure PostgresMultitenant.tenant_table_resolver.' unless resolver

  tenant_scoped_tables = resolver.call

  system_tables = PostgresMultitenant.configuration.system_tables

  # Combine and deduplicate
  @tenant_tables_cache = (Array(tenant_scoped_tables) + Array(system_tables)).uniq.sort
end

.terminate_schema_connections(name) ⇒ Object

Terminate all active connections to a specific schema

Parameters:

  • name (String) —

    The schema name



218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 218

def terminate_schema_connections(name)
  Rails.logger.info("Terminating known schema-scoped idle connections for schema: #{name}")

  sql = <<~SQL.squish
    SELECT pg_terminate_backend(pg_stat_activity.pid)
    FROM pg_stat_activity
    WHERE pg_stat_activity.datname = current_database()
      AND pid <> pg_backend_pid()
      AND state = 'idle'
      AND query ILIKE 'SET search_path TO "#{connection.quote_string(name)}"%'
  SQL

  begin
    connection.execute(sql)
  rescue StandardError => e
    Rails.logger.warn("Error terminating connections for schema #{name}: #{e.message}")
    # Don't raise - this is a best-effort operation
  end
end

.validate_tenant_table_configuration! ⇒ Object

Validate that tenant table configuration is correct Ensures models are loaded and schema matches expectations



32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 32

def self.validate_tenant_table_configuration!
  detected = tenant_tables

  # Ensure we found some tables
  if detected.empty?
    raise StandardError, 'No tenant tables detected. Configure PostgresMultitenant.tenant_table_resolver.'
  end

  # In test/development, validate that schema has tenant foreign key columns
  if Rails.env.test? || Rails.env.development?
    system_tables = PostgresMultitenant.configuration.system_tables
    tenant_foreign_key = PostgresMultitenant.configuration.tenant_foreign_key
    missing_org_id = detected.reject { |table| system_tables.include?(table) }
                             .select { |table| !column_exists_in_schema?(table, tenant_foreign_key) }

    if missing_org_id.any?
      Rails.logger.warn(
        "WARNING: These tenant tables are missing #{tenant_foreign_key} column: #{missing_org_id.join(', ')}"
      )
    end
  end

  detected
end

.with_schema(name, &block) ⇒ Object

Switch to a tenant schema for a block



88
89
90
91
92
93
94
95
96
97
98
99
100
101
 
# File 'lib/postgres_multitenant/schema_manager.rb', line 88

def with_schema(name, &block)
  raise ArgumentError, 'Block required' unless block

  name = normalize_schema_name(name)
  validate_schema_name!(name) unless name == 'public'

  original = current_schema
  switched = false
  switch_to(name)
  switched = true
  yield
ensure
  switch_to(original) if switched && original
end