Add temp_table_name option to with_temp_table (#32)

Gemfile

@@ -5,25 +5,24 @@ source "https://rubygems.org"
 # Specify your gem's dependencies in umbrellio_utils.gemspec
 gemspec
 
-group :development, :test do
-  gem "activesupport"
-  gem "bundler"
-  gem "ci-helper"
-  gem "http"
-  gem "nokogiri"
-  gem "nori"
-  gem "pg"
-  gem "pry"
-  gem "rake"
-  gem "rspec"
-  gem "rspec-json_matcher"
-  gem "rubocop-config-umbrellio"
-  gem "semantic_logger"
-  gem "sequel"
-  gem "sequel-batches"
-  gem "simplecov"
-  gem "simplecov-lcov"
-  gem "table_sync"
-  gem "timecop"
-  gem "yard"
-end
+gem "activesupport"
+gem "bundler"
+gem "ci-helper"
+gem "http"
+gem "net-pop", github: "ruby/net-pop" # See https://stackoverflow.com/questions/78617432/strange-bundle-update-issue-disappearing-net-pop-0-1-2-dependency
+gem "nokogiri"
+gem "nori"
+gem "pg"
+gem "pry"
+gem "rake"
+gem "rspec"
+gem "rspec-json_matcher"
+gem "rubocop-config-umbrellio"
+gem "semantic_logger"
+gem "sequel"
+gem "sequel-batches"
+gem "simplecov"
+gem "simplecov-lcov"
+gem "table_sync"
+gem "timecop"
+gem "yard"

Gemfile.lock

@@ -1,7 +1,14 @@
+GIT
+  remote: https://github.com/ruby/net-pop.git
+  revision: e8d0afe2773b9eb6a23c39e9e437f6fc0fc7c733
+  specs:
+    net-pop (0.1.2)
+      net-protocol
+
 PATH
   remote: .
   specs:
-    umbrellio-utils (1.3.0)
+    umbrellio-utils (1.4.0)
       memery (~> 1)
 
 GEM
@@ -106,7 +113,7 @@ GEM
     docile (1.4.0)
     domain_name (0.6.20240107)
     drb (2.2.1)
-    dry-inflector (1.0.0)
+    dry-inflector (1.1.0)
     erubi (1.13.0)
     ffi (1.17.0)
     ffi (1.17.0-arm64-darwin)
@@ -158,7 +165,6 @@ GEM
     net-imap (0.4.14)
       date
       net-protocol
-    net-pop (0.1.2)
     net-protocol (0.2.2)
       timeout
     net-smtp (0.5.0)
@@ -193,7 +199,7 @@ GEM
       sneakers (~> 2.0)
       tainbox
     racc (1.8.0)
-    rack (3.1.4)
+    rack (3.1.6)
     rack-session (2.0.0)
       rack (>= 3.0.0)
     rack-test (2.1.0)
@@ -304,7 +310,7 @@ GEM
     self_data (1.3.0)
     semantic_logger (4.15.0)
       concurrent-ruby (~> 1.0)
-    sequel (5.81.0)
+    sequel (5.82.0)
       bigdecimal
     sequel-batches (2.0.2)
       sequel
@@ -367,6 +373,7 @@ DEPENDENCIES
   bundler
   ci-helper
   http
+  net-pop!
   nokogiri
   nori
   pg

lib/umbrellio_utils/database.rb

@@ -22,20 +22,37 @@ def get_violated_constraint_name(exception)
       error.result.error_field(PG::Result::PG_DIAG_CONSTRAINT_NAME)
     end
 
-    def each_record(dataset, **options, &block)
-      primary_key = primary_key_from(dataset, **options)
+    def each_record(dataset, primary_key: nil, **options, &block)
+      primary_key = primary_key_from(dataset, primary_key: primary_key)
 
       with_temp_table(dataset, **options) do |ids|
         rows = ids.map { |id| row(id.is_a?(Hash) ? id.values : [id]) }
         dataset.model.where(row(primary_key) => rows).reverse(row(primary_key)).each(&block)
       end
     end
 
-    def with_temp_table(dataset, page_size: 1_000, sleep: nil, **options)
-      primary_key = primary_key_from(dataset, **options)
+    # Iterates over a dataset and yields batches of primary keys.
+    # First, a temporary table is created and populated with dataset primary keys.
+    # After that, a batch of rows is deleted from the temp table on each iteration
+    # and gets yielded to the caller.
+    # @option [Integer] page_size max size of each yielded PK batch
+    # @option [Integer] sleep interval to sleep between each iteration
+    # @option [Array] primary_key custom primary key to use for dataset
+    # @option [Symbol, String] temp_table_name custom name for temporary table,
+    #   table is reused if already exists
+    def with_temp_table(
+      dataset,
+      page_size: 1_000,
+      sleep: nil,
+      primary_key: nil,
+      temp_table_name: nil
+    )
+      primary_key = primary_key_from(dataset, primary_key: primary_key)
       sleep_interval = sleep_interval_from(sleep)
 
-      temp_table_name = create_temp_table(dataset, primary_key: primary_key)
+      temp_table_name = create_temp_table(
+        dataset, primary_key: primary_key, temp_table_name: temp_table_name&.to_sym
+      )
 
       pk_set = []
 
@@ -49,23 +66,26 @@ def with_temp_table(dataset, page_size: 1_000, sleep: nil, **options)
 
         Kernel.sleep(sleep_interval) if sleep_interval.positive?
       end
-    ensure
+
       DB.drop_table(temp_table_name)
     end
 
-    def create_temp_table(dataset, **options)
+    def create_temp_table(dataset, primary_key: nil, temp_table_name: nil)
       time = Time.current
       model = dataset.model
-      temp_table_name = :"temp_#{model.table_name}_#{time.to_i}_#{time.nsec}"
-      primary_key = primary_key_from(dataset, **options)
+
+      temp_table_name ||= :"temp_#{model.table_name}_#{time.to_i}_#{time.nsec}"
+      return temp_table_name if DB.table_exists?(temp_table_name)
+
+      primary_key = primary_key_from(dataset, primary_key: primary_key)
 
       DB.create_table(temp_table_name, unlogged: true) do
         primary_key.each do |field|
           type = model.db_schema[field][:db_type]
-          column field, type
+          column(field, type)
         end
 
-        primary_key primary_key
+        primary_key(primary_key)
       end
 
       insert_ds = dataset.select(*qualified_pk(model.table_name, primary_key))

lib/umbrellio_utils/version.rb

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

spec/umbrellio_utils/database_spec.rb

@@ -126,5 +126,22 @@
         end
       end
     end
+
+    context "custom temp table name provided" do
+      let(:options) { Hash[temp_table_name: "custom"] }
+
+      it "uses it" do
+        expect(DB).to receive(:create_table).with(:custom, unlogged: true).and_call_original
+        expect(result_emails).to eq(reversed_emails)
+      end
+    end
+
+    context "some invalid option provided" do
+      let(:options) { Hash[invalid: 1] }
+
+      it "raises error" do
+        expect { result_emails }.to raise_error(ArgumentError)
+      end
+    end
   end
 end