extract primary key from dataset if param is not present (#30)

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    umbrellio-utils (1.2.0)
+    umbrellio-utils (1.2.1)
       memery (~> 1)
 
 GEM

lib/umbrellio_utils/database.rb

@@ -5,6 +5,7 @@ module Database
     extend self
 
     HandledConstaintError = Class.new(StandardError)
+    InvalidPkError = Class.new(StandardError)
 
     def handle_constraint_error(constraint_name, &block)
       DB.transaction(savepoint: true, &block)
@@ -22,7 +23,7 @@ def get_violated_constraint_name(exception)
     end
 
     def each_record(dataset, **options, &block)
-      primary_key = primary_key_from(**options)
+      primary_key = primary_key_from(dataset, **options)
 
       with_temp_table(dataset, **options) do |ids|
         rows = ids.map { |id| row(id.is_a?(Hash) ? id.values : [id]) }
@@ -31,7 +32,7 @@ def each_record(dataset, **options, &block)
     end
 
     def with_temp_table(dataset, page_size: 1_000, sleep: nil, **options)
-      primary_key = primary_key_from(**options)
+      primary_key = primary_key_from(dataset, **options)
       sleep_interval = sleep_interval_from(sleep)
 
       temp_table_name = create_temp_table(dataset, primary_key: primary_key)
@@ -62,7 +63,7 @@ def create_temp_table(dataset, **options)
       time = Time.current
       model = dataset.model
       temp_table_name = "temp_#{model.table_name}_#{time.to_i}_#{time.nsec}".to_sym
-      primary_key = primary_key_from(**options)
+      primary_key = primary_key_from(dataset, **options)
 
       DB.create_table(temp_table_name, unlogged: true) do
         primary_key.each do |field|
@@ -85,8 +86,14 @@ def row(values)
       Sequel.function(:row, *values)
     end
 
-    def primary_key_from(**options)
-      Array(options.fetch(:primary_key, :id))
+    def extract_primary_key(dataset)
+      dataset.db.schema(dataset.first_source).select { |x| x[1][:primary_key] }.map(&:first)
+    end
+
+    def primary_key_from(dataset, **options)
+      Array(options[:primary_key] || extract_primary_key(dataset)).tap do |primary_key|
+        raise InvalidPkError if primary_key.empty?
+      end
     end
 
     def qualified_pk(table_name, primary_key)

lib/umbrellio_utils/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module UmbrellioUtils
-  VERSION = "1.2.0"
+  VERSION = "1.2.1"
 end

spec/umbrellio_utils/database_spec.rb

@@ -65,10 +65,12 @@
 
       let(:nicks) { complex_users_data.pluck(:nick) }
 
+      let(:primary_key_param) { %i[geo nick] }
+
       subject(:result_nicks) do
         users = []
 
-        described_class.each_record(ComplexUser.dataset, primary_key: %i[geo nick]) do |user|
+        described_class.each_record(ComplexUser.dataset, primary_key: primary_key_param) do |user|
           users << user
         end
 
@@ -79,6 +81,23 @@
         expect(result_nicks).to match_array(nicks)
         expect(sleep_calls).to eq([])
       end
+
+      context "without primary key param" do
+        let(:primary_key_param) { nil }
+
+        it "yields all records" do
+          expect(result_nicks).to match_array(nicks)
+          expect(sleep_calls).to eq([])
+        end
+      end
+
+      context "with invalid primary key param" do
+        let(:primary_key_param) { [] }
+
+        it "raises InvalidPkError" do
+          expect { result_nicks }.to raise_error(UmbrellioUtils::Database::InvalidPkError)
+        end
+      end
     end
 
     context "smaller page_size and numeric sleep value" do