Add enum plugin for treating columns as enums in a model

CHANGELOG

@@ -1,3 +1,7 @@
+=== master
+
+* Add enum plugin for treating columns as enums in a model (jeremyevans) (#1839)
+
 === 5.53.0 (2022-02-01)
 
 * Make Dataset#_sql_comment private when using the Database sql_comments extension (jeremyevans)

lib/sequel/plugins/enum.rb

@@ -0,0 +1,124 @@
+# frozen-string-literal: true
+
+module Sequel
+  module Plugins
+    # The enum plugin allows for easily adding methods to modify the value of
+    # a column.  It allows treating the column itself as an enum, returning a
+    # symbol for the related enum value.  It also allows for setting up dataset
+    # methods to easily find records having or not having each enum value.
+    #
+    # After loading the plugin, you can call the +enum+ method to define the
+    # methods.  The +enum+ method accepts a symbol for the underlying
+    # database column, and a hash with symbol keys for the enum values.
+    # For example, the following call:
+    #
+    #   Album.enum :status_id, good: 1, bad: 2
+    #
+    # Will define the following instance methods:
+    #
+    # Album#good! :: Change +status_id+ to +1+ (does not save the receiver)
+    # Album#bad! :: Change +status_id+ to +2+ (does not save the receiver)
+    # Album#good? :: Return whether +status_id+ is +1+
+    # Album#bad? :: Return whether +status_id+ is +2+
+    #
+    # It will override the following instance methods:
+    #
+    # Album#status_id :: Return +:good+/+:bad+ instead of +1+/+2+ (other values returned as-is)
+    # Album#status_id= :: Allow calling with +:good+/+:bad+ to set +status_id+ to +1+/+2+ (other values,
+    #                     such as <tt>'good'</tt>/<tt>'bad'</tt> set as-is)
+    #
+    # If will define the following dataset methods:
+    #
+    # Album.dataset.good :: Return a dataset filtered to rows where +status_id+ is +1+
+    # Album.dataset.not_good :: Return a dataset filtered to rows where +status_id+ is not +1+
+    # Album.dataset.bad:: Return a dataset filtered to rows where +status_id+ is +2+
+    # Album.dataset.not_bad:: Return a dataset filtered to rows where +status_id+ is not +2+
+    #
+    # When calling +enum+, you can also provide the following options:
+    #
+    # :prefix :: Use a prefix for methods defined for each enum value. If +true+ is provided at the value, use the column name as the prefix.
+    #            For example, with <tt>prefix: 'status'</tt>, the instance methods defined above would be +status_good?+, +status_bad?+,
+    #            +status_good!+, and +status_bad!+, and the dataset methods defined would be +status_good+, +status_not_good+, +status_bad+,
+    #            and +status_not_bad+.
+    # :suffix :: Use a suffix for methods defined for each enum value. If +true+ is provided at the value, use the column name as the suffix.
+    #            For example, with <tt>suffix: 'status'</tt>, the instance methods defined above would be +good_status?+, +bad_status?+,
+    #            +good_status!+, and +bad_status!+, and the dataset methods defined would be +good_status+, +not_good_status+, +bad_status+,
+    #            and +not_bad_status+.
+    # :override_accessors :: Set to +false+ to not override the column accessor methods.
+    # :dataset_methods :: Set to +false+ to not define dataset methods.
+    #
+    # Note that this does not use a true enum column in the database.  If you are
+    # looking for enum support in the database, and your are using PostgreSQL,
+    # Sequel supports that via the pg_enum Database extension.
+    #
+    # Usage:
+    #
+    #   # Make all model subclasses handle enums
+    #   Sequel::Model.plugin :enum
+    #
+    #   # Make the Album class handle enums
+    #   Album.plugin :enum
+    module Enum
+      module ClassMethods
+        # Define instance and dataset methods in this class to treat column
+        # as a enum.  See Enum documentation for usage.
+        def enum(column, values, opts=OPTS)
+          raise Sequel::Error, "enum column must be a symbol" unless column.is_a?(Symbol)
+          raise Sequel::Error, "enum values must be provided as a hash with symbol keys" unless values.is_a?(Hash) && values.all?{|k,| k.is_a?(Symbol)}
+
+          if prefix = opts[:prefix]
+            prefix = column if prefix == true
+            prefix = "#{prefix}_"
+          end
+
+          if suffix = opts[:suffix]
+            suffix = column if suffix == true
+            suffix = "_#{suffix}"
+          end
+
+          values = Hash[values].freeze
+          inverted = values.invert.freeze
+
+          unless @enum_methods
+            @enum_methods = Module.new
+            include @enum_methods
+          end
+
+          @enum_methods.module_eval do
+            unless opts[:override_accessors] == false
+              define_method(column) do
+                v = super()
+                inverted.fetch(v, v)
+              end
+
+              define_method(:"#{column}=") do |v|
+                super(values.fetch(v, v))
+              end
+            end
+
+            values.each do |key, value|
+              define_method(:"#{prefix}#{key}#{suffix}!") do
+                self[column] = value
+                nil
+              end
+
+              define_method(:"#{prefix}#{key}#{suffix}?") do
+                self[column] == value
+              end
+            end
+          end
+
+          unless opts[:dataset_methods] == false
+            dataset_module do
+              values.each do |key, value|
+                cond = Sequel[column=>value]
+                where :"#{prefix}#{key}#{suffix}", cond
+                where :"#{prefix}not_#{key}#{suffix}", ~cond
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

spec/extensions/enum_spec.rb

@@ -0,0 +1,180 @@
+require_relative "spec_helper"
+
+describe "Sequel enum plugin" do
+  before do
+    @Album = Class.new(Sequel::Model(Sequel.mock[:albums]))
+    @Album.columns :id, :status_id
+    @Album.plugin :enum
+    @Album.enum :status_id, :good=>3, :bad=>5
+    @album = @Album.load(:status_id=>3)
+  end
+
+  it "should add enum_value! and enum_value? methods for setting/checking the enum values" do
+    @album.good?.must_equal true
+    @album.bad?.must_equal false
+
+    @album.bad!.must_be_nil
+    @album.good?.must_equal false
+    @album.bad?.must_equal true
+
+    @album.good!.must_be_nil
+    @album.good?.must_equal true
+    @album.bad?.must_equal false
+  end
+
+  it "should have column method convert to enum value if possible" do
+    @album.status_id.must_equal :good
+    @album.bad!
+    @album.status_id.must_equal :bad
+    @album[:status_id] = 3
+    @album.status_id.must_equal :good
+  end
+
+  it "should have the column method pass non-enum values through" do
+    @album[:status_id] = 4
+    @album.status_id.must_equal 4
+  end
+
+  it "should have column= handle enum values" do
+    @album.status_id = :bad
+    @album[:status_id].must_equal 5
+    @album.good?.must_equal false
+    @album.bad?.must_equal true
+
+    @album.status_id = :good
+    @album[:status_id].must_equal 3
+    @album.good?.must_equal true
+    @album.bad?.must_equal false
+  end
+
+  it "should have column= handle non-enum values" do
+    @album.status_id = 5
+    @album[:status_id].must_equal 5
+    @album.good?.must_equal false
+    @album.bad?.must_equal true
+  end
+
+  it "should setup dataset methods for each value" do
+    ds = @Album.where(:id=>1)
+    ds.good.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id = 3))"
+    ds.not_good.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id != 3))"
+    ds.bad.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id = 5))"
+    ds.not_bad.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id != 5))"
+  end
+end
+
+describe "Sequel enum plugin" do
+  before do
+    @Album = Class.new(Sequel::Model(Sequel.mock[:albums]))
+    @Album.columns :id, :status_id
+    @Album.plugin :enum
+    @album = @Album.load(:status_id=>3)
+  end
+
+  it "should allow overriding methods in class and calling super" do
+    @Album.enum :status_id, {:good=>3, :bad=>5}, :override_accessors=>false
+    bad = nil
+    @Album.class_eval do
+      define_method(:bad?) do
+        bad.nil? ? super() : bad
+      end
+    end
+
+    @album.bad?.must_equal false
+    bad = true
+    @album.bad?.must_equal true
+    bad = false
+    @album.bad?.must_equal false
+    bad = nil
+    @album.bad!
+    @album.bad?.must_equal true
+  end
+
+  it "should not override accessor methods for each value if :override_accessors option is false" do
+    @Album.enum :status_id, {:good=>3, :bad=>5}, :override_accessors=>false
+    @album.status_id.must_equal 3
+    @album.status_id = :bad
+    @album.status_id.must_equal :bad
+    @album.bad!
+    @album.status_id.must_equal 5
+  end
+
+  it "should not setup dataset methods for each value if :dataset_methods option is false" do
+    @Album.enum :status_id, {:good=>3, :bad=>5}, :dataset_methods=>false
+    ds = @Album.where(:id=>1)
+    ds.wont_respond_to(:good)
+    ds.wont_respond_to(:not_good)
+    ds.wont_respond_to(:bad)
+    ds.wont_respond_to(:not_bad)
+  end
+
+  it "should handle :prefix=>true option" do
+    @Album.enum :status_id, {:good=>3, :bad=>5}, :prefix=>true
+
+    @album.status_id_good?.must_equal true
+    @album.status_id_bad!
+    @album.status_id_bad?.must_equal true
+
+    ds = @Album.where(:id=>1)
+    ds.status_id_good.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id = 3))"
+    ds.status_id_not_good.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id != 3))"
+  end
+
+  it "should handle :prefix=>string option" do
+    @Album.enum :status_id, {:good=>3, :bad=>5}, :prefix=>'status'
+
+    @album.status_good?.must_equal true
+    @album.status_bad!
+    @album.status_bad?.must_equal true
+
+    ds = @Album.where(:id=>1)
+    ds.status_good.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id = 3))"
+    ds.status_not_good.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id != 3))"
+  end
+
+  it "should handle :suffix=>true option" do
+    @Album.enum :status_id, {:good=>3, :bad=>5}, :suffix=>true
+
+    @album.good_status_id?.must_equal true
+    @album.bad_status_id!
+    @album.bad_status_id?.must_equal true
+
+    ds = @Album.where(:id=>1)
+    ds.good_status_id.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id = 3))"
+    ds.not_good_status_id.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id != 3))"
+  end
+
+  it "should handle :suffix=>true option" do
+    @Album.enum :status_id, {:good=>3, :bad=>5}, :suffix=>'status'
+
+    @album.good_status?.must_equal true
+    @album.bad_status!
+    @album.bad_status?.must_equal true
+
+    ds = @Album.where(:id=>1)
+    ds.good_status.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id = 3))"
+    ds.not_good_status.sql.must_equal "SELECT * FROM albums WHERE ((id = 1) AND (status_id != 3))"
+  end
+
+  it "should support multiple emums per class" do 
+    @Album.enum :id, {:odd=>1, :even=>2}
+    @Album.enum :status_id, {:good=>3, :bad=>5}
+    @album = @Album.load(:id=>1, :status_id=>3)
+    @album.odd?.must_equal true
+    @album.even?.must_equal false
+    @album.good?.must_equal true
+    @album.bad?.must_equal false
+  end
+
+  it "raises Error for column that isn't a symbol" do
+    proc{@Album.enum 'status_id', :good=>3, :bad=>5}.must_raise Sequel::Error
+  end
+
+  it "raises Error for non-hash values" do
+    proc{@Album.enum :status_id, [:good, :bad]}.must_raise Sequel::Error
+  end
+
+  it "raises Error for values hash with non-symbol keys" do
+    proc{@Album.enum :status_id, 'good'=>3, :bad=>5}.must_raise Sequel::Error
+  end
+end

www/pages/plugins.html.erb

@@ -34,6 +34,7 @@
 <li><a href="rdoc-plugins/classes/Sequel/Plugins/ColumnEncryption.html">column_encryption</a>: Encrypt column values stored in the database.</li>
 <li><a href="rdoc-plugins/classes/Sequel/Plugins/Dirty.html">dirty</a>: Allows you get get initial values of columns after changing the values.</li>
 <li><a href="rdoc-plugins/classes/Sequel/Plugins/DefaultsSetter.html">defaults_setter</a>: Get default values for new models before saving.</li>
+<li><a href="rdoc-plugins/classes/Sequel/Plugins/Enum.html">enum</a>: Allows for treating a column as a enum.</li>
 <li><a href="rdoc-plugins/classes/Sequel/Plugins/ForceEncoding.html">force_encoding</a>: Forces the all model column string values to a given encoding.</li>
 <li><a href="rdoc-plugins/classes/Sequel/Plugins/InputTransformer.html">input_transformer</a>: Automatically transform input to model column setters.</li>
 <li><a href="rdoc-plugins/classes/Sequel/Plugins/LazyAttributes.html">lazy_attributes</a>: Allows you to set some attributes that should not be loaded by default, but only loaded when an object requests them.</li>