DDL Improvements

README.md

@@ -54,14 +54,17 @@ When a migration modifies a table that exists in `public` but not yet in the bra
 ```
 add_column :users, :bio, :string
 
-1. CREATE TABLE branch_feature_payments.users (LIKE public.users INCLUDING ALL)
+1. CREATE TABLE branch_feature_payments.users (LIKE public.users ...)
 2. INSERT INTO branch_feature_payments.users SELECT * FROM public.users
-3. ALTER TABLE users ADD COLUMN bio VARCHAR
+3. Rebuild indexes, foreign keys, exclusion and unique constraints
+4. ALTER TABLE users ADD COLUMN bio VARCHAR
 ```
 
-The public table is never touched. Step 3 operates on the shadow because `search_path` resolves `users` to the branch schema first.
+The public table is never touched. Step 4 operates on the shadow because `search_path` resolves `users` to the branch schema first.
 
-This applies to any DDL that modifies an existing table: `add_column`, `remove_column`, `add_index`, `add_foreign_key`, `drop_table`, and so on.
+The adapter intercepts every `SchemaStatements` method that takes a table name — column DDL, index DDL, constraint DDL, `drop_table`, `change_table`, timestamps, references, and so on. Enum type DDL (`drop_enum`, `rename_enum`, `add_enum_value`, `rename_enum_value`) is also intercepted so enum modifications stay on the branch.
+
+When you `drop_table` on a branch, a lightweight tombstone is created in a separate schema to block `search_path` fallthrough to public. The public table is untouched and the schema dumper excludes it.
 
 ### schema.rb stays clean
 
@@ -155,5 +158,6 @@ All standard PostgreSQL connection parameters work as normal (`host`, `port`, `u
 
 - **Rails + PostgreSQL only** — uses PostgreSQL schemas and `search_path`
 - **Development and test only** — production should use the standard `postgresql` adapter
-- **Non-ActiveRecord DDL** — raw SQL outside of migrations bypasses the shadow rule
-- **Sequences on shadowed tables** — `rename_table` on a shadowed table with serial columns works, but the sequence keeps its original name
+- **Raw SQL DDL** — `execute("CREATE TRIGGER ...")` and similar raw SQL bypasses the shadow rule. If another intercepted DDL method has already shadowed the table in the same migration, the raw SQL will land on the branch copy via `search_path`. If the raw SQL is the first thing touching that table, it will hit `public`. Structure your migration so that an intercepted method (e.g. `add_column`) runs first, or use `db:branch:reset && db:migrate` to rebuild the branch from scratch.
+- **Sequences on shadowed tables** — `rename_table` on a shadowed table with serial columns works, but the sequence keeps its original name since renaming it would affect all branches
+- **Concurrent shadow creation** — if parallel workers (parallel tests, multiple consoles) shadow the same table simultaneously, a race condition can cause duplicate data or a `DuplicateTable` error. Restart the workers or run `db:branch:reset` to recover

lib/active_record/connection_adapters/postgresql/branched/adapter.rb

@@ -2,31 +2,36 @@ module ActiveRecord
   module ConnectionAdapters
     module PostgreSQL
       module Branched
+        # All SchemaStatements methods whose first parameter name
+        # contains "table" — covers table_name, from_table, table_1,
+        # table_names, and any future variant.
+        def self.table_methods
+          [
+            ActiveRecord::ConnectionAdapters::SchemaStatements,
+            ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaStatements
+          ].flat_map { |mod|
+            mod.instance_methods(false).select { |m|
+              first_param = mod.instance_method(m).parameters.first
+              first_param && first_param.last.to_s.include?("table")
+            }
+          }.uniq
+        end
+
         class Adapter < ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
           ADAPTER_NAME = "PostgreSQL Branched"
 
-          SHADOW_BEFORE = %i[
-            add_column
-            remove_column
-            rename_column
-            change_column
-            change_column_default
-            change_column_null
-            change_column_comment
-            change_table_comment
-            add_index
-            remove_index
-            rename_index
-            add_foreign_key
-            remove_foreign_key
-            add_check_constraint
-            remove_check_constraint
-            validate_foreign_key
-            validate_check_constraint
-            drop_table
-            change_table
-            bulk_change_table
-          ].freeze
+          # Methods where the first argument is not the table to shadow,
+          # or where shadowing would be wasteful (join table methods pass
+          # table_1, not the derived join table name).
+          SHADOW_SKIP = %i[create_join_table drop_join_table].freeze
+
+          # Methods with custom shadow handling outside the generic loop.
+          SHADOW_HANDLED = %i[rename_table].freeze
+
+          # Everything else with a table param gets the generic wrapper.
+          # Shadow#call is idempotent — if the table doesn't exist in
+          # public or is already shadowed, it's a no-op.
+          SHADOW_BEFORE = (Branched.table_methods - SHADOW_SKIP - SHADOW_HANDLED).freeze
 
           def initialize(...)
             super
@@ -40,9 +45,38 @@ def configure_connection
           end
 
           SHADOW_BEFORE.each do |method|
-            define_method(method) do |table_name, *args, **kwargs, &block|
-              @shadow&.call(table_name)
-              super(table_name, *args, **kwargs, &block)
+            define_method(method) do |*args, **kwargs, &block|
+              @shadow&.call(args.first)
+              super(*args, **kwargs, &block)
+            end
+          end
+
+          # Shadow each table so DROP resolves to the branch copy, then
+          # create a tombstone in the dropped schema to block fallthrough.
+          def drop_table(*table_names, **options)
+            if @branch_manager.primary_branch?
+              super
+            else
+              table_names.each { |t| @shadow&.call(t) }
+              super
+              table_names.each { |t| @shadow&.drop_table(t) }
+            end
+          end
+
+          # create_table after drop_table should work — remove the
+          # tombstone so the new table takes precedence.
+          def create_table(*args, **kwargs, &block)
+            @shadow&.undrop_table(args.first)
+            super
+          end
+
+          # Enum DDL resolves unqualified type names through search_path,
+          # so modifications on a feature branch would otherwise mutate
+          # public. Shadow the enum into the branch schema first.
+          %i[drop_enum rename_enum add_enum_value rename_enum_value].each do |method|
+            define_method(method) do |name, *args, **kwargs, &block|
+              @shadow&.shadow_enum(name)
+              super(name, *args, **kwargs, &block)
             end
           end
 
@@ -59,6 +93,10 @@ def rename_table(table_name, new_name, **options)
           end
 
           attr_reader :branch_manager
+
+          def table_dropped_on_branch?(table_name)
+            @shadow&.dropped?(table_name) || false
+          end
         end
       end
     end

lib/active_record/connection_adapters/postgresql/branched/branch_manager.rb

@@ -133,7 +133,8 @@ def drop_schema
           end
 
           def set_search_path
-            @connection.schema_search_path = "#{@branch_schema}, public"
+            dropped_schema = "#{@branch_schema}#{Shadow::DROPPED_SUFFIX}"
+            @connection.schema_search_path = "#{dropped_schema}, #{@branch_schema}, public"
           end
 
           def shadow_migration_tables(shadow)

lib/active_record/connection_adapters/postgresql/branched/schema_dumper.rb

@@ -27,6 +27,7 @@ def tables(stream)
 
             table_names = @connection.tables.uniq.sort
             table_names.reject! { |t| ignored?(t) }
+            table_names.reject! { |t| @connection.table_dropped_on_branch?(t) }
 
             table_names.each_with_index do |table_name, index|
               table(table_name, stream)

lib/active_record/connection_adapters/postgresql/branched/shadow.rb

@@ -3,21 +3,98 @@ module ConnectionAdapters
     module PostgreSQL
       module Branched
         class Shadow
+          DROPPED_SUFFIX = "__dropped"
+
           def initialize(connection, branch_schema)
             @connection = connection
             @branch_schema = branch_schema
+            @dropped_schema = "#{branch_schema}#{DROPPED_SUFFIX}"
           end
 
+          attr_reader :dropped_schema
+
           def call(table_name)
             table = table_name.to_s
             return unless exists_in_public?(table)
             return if already_shadowed?(table)
+            return if dropped?(table)
 
             create_shadow(table)
           end
 
+          def shadow_enum(name)
+            name = name.to_s
+            return unless enum_exists_in_public?(name)
+            return if enum_exists_in_branch?(name)
+
+            labels = @connection.select_values(<<~SQL)
+              SELECT e.enumlabel FROM pg_enum e
+                JOIN pg_type t ON t.oid = e.enumtypid
+                JOIN pg_namespace n ON n.oid = t.typnamespace
+              WHERE n.nspname = 'public' AND t.typname = #{@connection.quote(name)}
+              ORDER BY e.enumsortorder
+            SQL
+
+            quoted_branch = @connection.quote_column_name(@branch_schema)
+            quoted_labels = labels.map { |l| @connection.quote(l) }.join(", ")
+            @connection.execute(<<~SQL)
+              CREATE TYPE #{quoted_branch}.#{@connection.quote_column_name(name)}
+                AS ENUM (#{quoted_labels})
+            SQL
+          end
+
+          def drop_table(table)
+            table = table.to_s
+            return unless exists_in_public?(table)
+
+            ensure_dropped_schema
+            quoted_dropped = @connection.quote_column_name(@dropped_schema)
+            quoted_table = @connection.quote_column_name(table)
+            @connection.execute(<<~SQL)
+              CREATE TABLE #{quoted_dropped}.#{quoted_table}
+                (LIKE public.#{quoted_table})
+            SQL
+          end
+
+          def undrop_table(table)
+            table = table.to_s
+            return unless dropped?(table)
+            quoted_dropped = @connection.quote_column_name(@dropped_schema)
+            quoted_table = @connection.quote_column_name(table)
+            @connection.execute("DROP TABLE #{quoted_dropped}.#{quoted_table}")
+          end
+
+          def dropped?(table_name)
+            @connection.select_value(<<~SQL) == 1
+              SELECT 1 FROM information_schema.tables
+              WHERE table_schema = #{@connection.quote(@dropped_schema)}
+                AND table_name = #{@connection.quote(table_name.to_s)}
+            SQL
+          end
+
           private
 
+          def enum_exists_in_public?(name)
+            @connection.select_value(<<~SQL) == 1
+              SELECT 1 FROM pg_type t
+                JOIN pg_namespace n ON n.oid = t.typnamespace
+              WHERE n.nspname = 'public' AND t.typname = #{@connection.quote(name)} AND t.typtype = 'e'
+            SQL
+          end
+
+          def enum_exists_in_branch?(name)
+            @connection.select_value(<<~SQL) == 1
+              SELECT 1 FROM pg_type t
+                JOIN pg_namespace n ON n.oid = t.typnamespace
+              WHERE n.nspname = #{@connection.quote(@branch_schema)} AND t.typname = #{@connection.quote(name)} AND t.typtype = 'e'
+            SQL
+          end
+
+          def ensure_dropped_schema
+            quoted = @connection.quote_column_name(@dropped_schema)
+            @connection.execute("CREATE SCHEMA IF NOT EXISTS #{quoted}")
+          end
+
           def exists_in_public?(table)
             @connection.select_value(<<~SQL) == 1
               SELECT 1 FROM information_schema.tables
@@ -40,14 +117,11 @@ def create_shadow(table)
             clone_structure(quoted_branch, quoted_table)
             copy_data(quoted_branch, quoted_table)
             clone_indexes(table, quoted_branch, quoted_table)
+            clone_constraints(table, quoted_branch, quoted_table)
           end
 
-          # Per-row index maintenance during INSERT SELECT dominates the cost
-          # of shadowing large tables. LIKE ... INCLUDING ALL builds every
-          # index up front, then the data copy updates each index for every
-          # row. EXCLUDING INDEXES lets the bulk copy run without index
-          # overhead; clone_indexes rebuilds them afterwards in a single
-          # sort pass per index, the standard Postgres bulk-load pattern.
+          # Structure without indexes — they're rebuilt after the bulk
+          # data copy for better performance (sort-based vs per-row).
           def clone_structure(quoted_branch, quoted_table)
             @connection.execute(<<~SQL)
               CREATE TABLE #{quoted_branch}.#{quoted_table}
@@ -76,10 +150,11 @@ def clone_indexes(table, quoted_branch, quoted_table)
             end
           end
 
-          # Redirect pg_get_indexdef output from public to the branch schema.
-          # The REPLACE uses the same quote_ident() that pg_get_indexdef uses
-          # to emit identifiers, so the anchor always matches regardless of
-          # whether the table name needs quoting.
+          # LIKE ... INCLUDING ALL EXCLUDING INDEXES drops all index-backed
+          # constraints (PK, unique, exclusion). clone_constraints re-emits
+          # unique and exclusion constraints which implicitly create their
+          # backing indexes, so we skip those here to avoid name collisions.
+          # PK indexes are rebuilt and reattached via USING INDEX.
           def index_metadata(table)
             @connection.select_rows(<<~SQL)
               SELECT i.relname, ix.indisprimary,
@@ -94,6 +169,40 @@ def index_metadata(table)
                 JOIN pg_namespace n ON n.oid = t.relnamespace
               WHERE n.nspname = 'public'
                 AND t.relname = #{@connection.quote(table)}
+                AND NOT EXISTS (
+                  SELECT 1 FROM pg_constraint c
+                  WHERE c.conindid = ix.indexrelid
+                    AND c.contype IN ('x', 'u')
+                )
+            SQL
+          end
+
+          # LIKE ... INCLUDING CONSTRAINTS copies check constraints and NOT
+          # NULL but never foreign keys (Postgres design). EXCLUDING INDEXES
+          # also drops exclusion and unique constraints since they're index-
+          # backed. Re-emit each from pg_constraint; the definitions from
+          # pg_get_constraintdef are self-contained and resolve correctly
+          # via the branch search_path without rewriting.
+          def clone_constraints(table, quoted_branch, quoted_table)
+            constraint_definitions(table).each do |conname, condef|
+              quoted_con = @connection.quote_column_name(conname)
+              @connection.execute(<<~SQL)
+                ALTER TABLE #{quoted_branch}.#{quoted_table}
+                  ADD CONSTRAINT #{quoted_con} #{condef}
+              SQL
+            end
+          end
+
+          def constraint_definitions(table)
+            @connection.select_rows(<<~SQL)
+              SELECT c.conname, pg_get_constraintdef(c.oid, true)
+              FROM pg_constraint c
+                JOIN pg_class t ON t.oid = c.conrelid
+                JOIN pg_namespace n ON n.oid = t.relnamespace
+              WHERE n.nspname = 'public'
+                AND t.relname = #{@connection.quote(table)}
+                AND c.contype IN ('f', 'x', 'u')
+              ORDER BY c.conname
             SQL
           end
         end

test/adapter_test.rb

@@ -53,7 +53,7 @@ def test_feature_branch_creates_schema
   def test_feature_branch_sets_search_path
     conn = connect(branch: "feature/payments")
     search_path = conn.select_value("SHOW search_path")
-    assert_equal "branch_feature_payments, public", search_path
+    assert_equal "branch_feature_payments__dropped, branch_feature_payments, public", search_path
   end
 
   def test_feature_branch_table_takes_precedence_over_public