From 14e859c6dd0451fb16eca24c47ab082cf06e7fdb Mon Sep 17 00:00:00 2001
From: Michael Herold <opensource@michaeljherold.com>
Date: Fri, 20 Oct 2023 22:54:56 -0500
Subject: [PATCH] Fix add_column behavior for missing type

When using `add_column`, the definition methods for `ksuid` and
`ksuid_binary` are not used like when using `table.ksuid` since
`add_column` creates the column directly.

This change makes it so any column definition can handle both `ksuid`
and `ksuid_binary` correctly by monkey-patching the method that creates
the column definitions.
---
 activerecord-ksuid/CHANGELOG.md               |  6 +++++
 .../lib/active_record/ksuid/railtie.rb        |  8 ++++---
 .../active_record/ksuid/table_definition.rb   | 22 +++++++++++++++++++
 .../spec/active_record/ksuid/railtie_spec.rb  | 17 ++++++++++++++
 4 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/activerecord-ksuid/CHANGELOG.md b/activerecord-ksuid/CHANGELOG.md
index b1e3563..8153234 100644
--- a/activerecord-ksuid/CHANGELOG.md
+++ b/activerecord-ksuid/CHANGELOG.md
@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
 
+## [Unreleased](https://github.com/michaelherold/ksuid-ruby/compare/v1.0.0...main)
+
+### Fixed
+
+- Using a `ksuid` or `ksuid_binary` type with `add_column` in migrations will now work correctly without throwing an error about an unknown column type.
+
 ## [1.0.0](https://github.com/michaelherold/ksuid-ruby/compare/v0.5.0...v1.0.0) - 2023-02-25
 
 ### Added
diff --git a/activerecord-ksuid/lib/active_record/ksuid/railtie.rb b/activerecord-ksuid/lib/active_record/ksuid/railtie.rb
index 3e95a6b..bd2a429 100644
--- a/activerecord-ksuid/lib/active_record/ksuid/railtie.rb
+++ b/activerecord-ksuid/lib/active_record/ksuid/railtie.rb
@@ -18,9 +18,11 @@ class Railtie < ::Rails::Railtie
         ActiveSupport.on_load :active_record do
           require 'active_record/ksuid/table_definition'
 
-          ActiveRecord::ConnectionAdapters::TableDefinition.include(
-            ActiveRecord::KSUID::TableDefinition
-          )
+          ActiveRecord::ConnectionAdapters::TableDefinition.descendants.each do |defn|
+            next unless defn.name
+
+            defn.prepend(ActiveRecord::KSUID::TableDefinition)
+          end
         end
       end
     end
diff --git a/activerecord-ksuid/lib/active_record/ksuid/table_definition.rb b/activerecord-ksuid/lib/active_record/ksuid/table_definition.rb
index fda031c..347ba7e 100644
--- a/activerecord-ksuid/lib/active_record/ksuid/table_definition.rb
+++ b/activerecord-ksuid/lib/active_record/ksuid/table_definition.rb
@@ -52,6 +52,28 @@ def ksuid(*args, **options)
       def ksuid_binary(*args, **options)
         args.each { |name| column(name, :binary, **options.merge(limit: 20)) }
       end
+
+      # Monkey-patches defining a new column within a table
+      #
+      # @api private
+      # @private
+      #
+      # @param name [String, Symbol] the name of the column
+      # @param type [String, Symbol] the type of the column
+      # @param options [Hash<Symbol, Object>] options for the definition
+      # @return [ActiveRecord::ConnectionAdapters::ColumnDefinition]
+      def new_column_definition(name, type, **options)
+        case type.to_s
+        when 'ksuid'
+          prefix_length = options.delete(:prefix)&.length || 0
+
+          super(name, :string, **options.merge(limit: 27 + prefix_length))
+        when 'ksuid_binary'
+          super(name, :binary, **options.merge(limit: 20))
+        else
+          super
+        end
+      end
     end
   end
 end
diff --git a/activerecord-ksuid/spec/active_record/ksuid/railtie_spec.rb b/activerecord-ksuid/spec/active_record/ksuid/railtie_spec.rb
index a0f9990..bea2ef3 100644
--- a/activerecord-ksuid/spec/active_record/ksuid/railtie_spec.rb
+++ b/activerecord-ksuid/spec/active_record/ksuid/railtie_spec.rb
@@ -257,6 +257,23 @@ class EventPrefix < ActiveRecord::Base
     end
   end
 
+  context 'when writing a migration that adds KSUID fields' do
+    it 'can use the ksuid and ksuid_binary field types' do
+      connection = ActiveRecord::Base.connection
+
+      connection.create_table :field_tests do |t|
+        t.string :name
+      end
+
+      expect do
+        connection.add_column :field_tests, :id_one, :ksuid
+        connection.add_column :field_tests, :id_two, :ksuid_binary
+      end.not_to raise_error
+    ensure
+      connection.drop_table :field_tests, if_exists: true
+    end
+  end
+
   matcher :issue_sql_queries do |expected|
     supports_block_expectations