Composite primary key

.ruby-version

@@ -1 +1 @@
-2.3.1
+2.5.3

lib/sequel/plugins/bitemporal.rb

@@ -31,19 +31,25 @@ def self.now
         Thread.current[THREAD_NOW_KEY] || DateTime.now
       end
 
-      def self.bitemporal_version_columns
-        @bitemporal_version_columns ||= [:master_id, :valid_from, :valid_to, :created_at, :expired_at]
+      def self.version_foreign_keys(master = nil)
+        return :master_id unless master
+        primary_key = [*master.primary_key]
+        primary_key.size > 1 ? primary_key : :master_id
       end
 
-      def self.bitemporal_excluded_columns
-        @bitemporal_excluded_columns ||= [:id, *bitemporal_version_columns]
+      def self.bitemporal_version_columns(master = nil)
+        [*version_foreign_keys(master), :valid_from, :valid_to, :created_at, :expired_at]
+      end
+
+      def self.bitemporal_excluded_columns(master = nil)
+        [:id, *bitemporal_version_columns(master)]
       end
 
       def self.configure(master, opts = {})
         version = opts[:version_class]
         raise Error, "please specify version class to use for bitemporal plugin" unless version
         return version.db.log_info("Version table does not exist for #{version.name}") unless version.db.table_exists?(version.table_name)
-        missing = bitemporal_version_columns - version.columns
+        missing = bitemporal_version_columns(master) - version.columns
         raise Error, "bitemporal plugin requires the following missing column#{"s" if missing.size>1} on version class: #{missing.join(", ")}" unless missing.empty?
 
         if Sequel::Plugins::Bitemporal.jdbc?(master.db)
@@ -76,7 +82,7 @@ def self.configure(master, opts = {})
           @audit_updated_by_method = opts.fetch(:audit_updated_by_method){ :updated_by }
           @propagate_per_column = opts.fetch(:propagate_per_column, false)
           @version_uses_string_nilifier = version.plugins.map(&:to_s).include? "Sequel::Plugins::StringNilifier"
-          @excluded_columns = Sequel::Plugins::Bitemporal.bitemporal_excluded_columns
+          @excluded_columns = Sequel::Plugins::Bitemporal.bitemporal_excluded_columns(master)
           @excluded_columns += Array opts[:excluded_columns] if opts[:excluded_columns]
           @use_ranges = if opts[:ranges]
             db = self.db
@@ -100,8 +106,8 @@ def self.current_versions_dataset
             end
           end
         end
-        master.one_to_many :versions, class: version, key: :master_id, graph_alias_base: master.versions_alias
-        master.one_to_one :current_version, class: version, key: :master_id, graph_alias_base: master.current_version_alias, :graph_block=>(proc do |j, lj, js|
+        master.one_to_many :versions, class: version, key: version_foreign_keys(master), graph_alias_base: master.versions_alias
+        master.one_to_one :current_version, class: version, key: version_foreign_keys(master), graph_alias_base: master.current_version_alias, :graph_block=>(proc do |j, lj, js|
           t = Sequel.delay{ ::Sequel::Plugins::Bitemporal.point_in_time }
           n = Sequel.delay{ ::Sequel::Plugins::Bitemporal.now }
           if master.use_ranges
@@ -135,7 +141,7 @@ def self.current_versions_dataset
             )
           )
         end
-        master.one_to_many :current_or_future_versions, class: version, key: :master_id, :graph_block=>(proc do |j, lj, js|
+        master.one_to_many :current_or_future_versions, class: version, key: version_foreign_keys(master), :graph_block=>(proc do |j, lj, js|
           t = Sequel.delay{ ::Sequel::Plugins::Bitemporal.point_in_time }
           n = Sequel.delay{ ::Sequel::Plugins::Bitemporal.now }
           if master.use_ranges
@@ -169,7 +175,7 @@ def self.current_versions_dataset
             Sequel.negate(Sequel.qualify(:current_or_future_versions, :id) => nil)
           )
         end
-        version.many_to_one :master, class: master, key: :master_id
+        version.many_to_one :master, class: master, key: version_foreign_keys(master)
         version.class_eval do
           if Sequel::Plugins::Bitemporal.jdbc?(master.db)
             plugin :typecast_on_load, *columns
@@ -319,7 +325,13 @@ def attributes
 
         def attributes=(attributes)
           @pending_version ||= begin
-            current_attributes = {master_id: id}
+            current_attributes =
+              if composite_primary_key?
+                version_values
+              else
+                { master_id: id }
+              end
+
             current_version.keys.each do |key|
               next if excluded_columns.include? key
               current_attributes[key] = current_version.send key
@@ -329,6 +341,10 @@ def attributes=(attributes)
           pending_version.set_all attributes
         end
 
+        def version_values
+          version_foreign_keys.map { |k| [k, public_send(k)] }.to_h
+        end
+
         def update_attributes(attributes={})
           self.attributes = attributes
           if save raise_on_failure: false
@@ -364,6 +380,10 @@ def after_save
           _refresh_set_values @values
         end
 
+        def composite_primary_key?
+          [*primary_key].size > 1
+        end
+
         def destroy
           point_in_time = ::Sequel::Plugins::Bitemporal.point_in_time
           versions_dataset.where(
@@ -471,6 +491,10 @@ def propagated_during_last_save
           @propagated_during_last_save ||= []
         end
 
+        def version_foreign_keys
+          composite_primary_key? ? primary_key : :master_id
+        end
+
       private
 
         def prepare_pending_version
@@ -608,7 +632,7 @@ def pending_version_holds_changes?
           columns.detect do |column|
             new_value = pending_version.send column
             case column
-            when :id, :master_id, :created_at, :expired_at
+            when :id, :created_at, :expired_at, *version_foreign_keys
               false
             when :valid_from
               pending_version.values.has_key?(:valid_from) && (

spec/composite_primary_key_spec.rb

@@ -0,0 +1,61 @@
+require "spec_helper"
+
+BT = Sequel::Plugins::Bitemporal
+
+describe "composite_primary_key" do
+  before :all do
+    setup_composite_primary_key
+  end
+
+  describe "::version_foreign_keys" do
+    context "when is invoked without args" do
+      it "returns master_id" do
+        expect(BT.bitemporal_excluded_columns).to(
+          eq([:id, :master_id, :valid_from, :valid_to, :created_at, :expired_at])
+        )
+      end
+    end
+  end
+
+  describe "missing required foreign keys in version table" do
+    ERROR_TEXT = "bitemporal plugin requires the following missing columns on version class: department_id, team_id"
+
+    it "raises error" do
+      expect do
+        setup_composite_primary_key(with_foreign_key: false)
+      end.to raise_error Sequel::Error, ERROR_TEXT
+    end
+  end
+
+  describe "versions work correctly" do
+    let(:master) { @master_class.new(name: "john Smith") }
+    let(:junior) { "Junior Ruby-developer" }
+    let(:senior) { "Senior Ruby-developer" }
+    let(:two_years_in_seconds) { 2 * 365 * 24 * 60 * 60 }
+
+    before do
+      setup_composite_primary_key
+      master.department_id = 1
+      master.team_id = 1
+      master.save
+
+      BT.at(Time.now) do
+        master.update_attributes(position: junior)
+      end
+
+      BT.at(Time.now + two_years_in_seconds) do
+        master.update_attributes(position: senior)
+      end
+    end
+
+    specify do
+      BT.at(Time.now) do
+        expect(master.reload.position).to eq(junior)
+      end
+
+      BT.at(Time.now + two_years_in_seconds) do
+        expect(master.reload.position).to eq(senior)
+      end
+    end
+  end
+end

spec/support/db.rb

@@ -63,6 +63,48 @@ def validate
     end
   end
 
+  def setup_composite_primary_key(with_foreign_key: true)
+    DB.drop_table(:employee_versions) if DB.table_exists?(:employee_versions)
+    DB.drop_table(:employees) if DB.table_exists?(:employees)
+
+    DB.create_table! :employees do
+      Integer :department_id
+      Integer :team_id
+      primary_key [:department_id, :team_id]
+      String :name, null: false
+    end
+
+    DB.create_table! :employee_versions do
+      if with_foreign_key
+        Integer :department_id
+        Integer :team_id
+
+        foreign_key [:department_id, :team_id], :employees
+      end
+
+      primary_key :id
+      String :position
+      Date :created_at
+      Date :expired_at
+      Date :valid_from
+      Date :valid_to
+    end
+
+    @version_class = Class.new Sequel::Model do
+      set_dataset :employee_versions
+    end
+
+    bitemporal_options = {
+      version_class: @version_class,
+      key: [:department_id, :team_id]
+    }
+
+    @master_class = Class.new Sequel::Model do
+      set_dataset :employees
+      plugin :bitemporal, bitemporal_options
+    end
+  end
+
   def db_truncate
     if DbHelpers.pg?
       @version_class.truncate cascade: true