From c2480db64d8dd4997d6cd1b915b215967c6f6db7 Mon Sep 17 00:00:00 2001
From: Matthew Riddle <mriddle89@gmail.com>
Date: Tue, 7 Apr 2020 14:25:13 +0200
Subject: [PATCH] Replace `with_connections` with `with_servers`

Refactoring so we're just working with one list of things.
---
 lib/global_uid/base.rb      | 24 +++++++--------
 test/lib/global_uid_test.rb | 58 ++++++++++++++++++-------------------
 2 files changed, 40 insertions(+), 42 deletions(-)

diff --git a/lib/global_uid/base.rb b/lib/global_uid/base.rb
index e3bfa39..8b28e3e 100644
--- a/lib/global_uid/base.rb
+++ b/lib/global_uid/base.rb
@@ -22,8 +22,8 @@ def self.create_uid_tables(id_table_name, options={})
 
       engine_stmt = "ENGINE=#{GlobalUid.configuration.storage_engine}"
 
-      with_connections do |connection|
-        connection.execute("CREATE TABLE IF NOT EXISTS `#{id_table_name}` (
+      with_servers do |server|
+        server.connection.execute("CREATE TABLE IF NOT EXISTS `#{id_table_name}` (
         `id` #{type} NOT NULL AUTO_INCREMENT,
         `stub` char(1) NOT NULL DEFAULT '',
         PRIMARY KEY (`id`),
@@ -31,13 +31,13 @@ def self.create_uid_tables(id_table_name, options={})
         ) #{engine_stmt}")
 
         # prime the pump on each server
-        connection.execute("INSERT IGNORE INTO `#{id_table_name}` VALUES(#{start_id}, 'a')")
+        server.connection.execute("INSERT IGNORE INTO `#{id_table_name}` VALUES(#{start_id}, 'a')")
       end
     end
 
     def self.drop_uid_tables(id_table_name)
-      with_connections do |connection|
-        connection.execute("DROP TABLE IF EXISTS `#{id_table_name}`")
+      with_servers do |server|
+        server.connection.execute("DROP TABLE IF EXISTS `#{id_table_name}`")
       end
     end
 
@@ -67,7 +67,7 @@ def self.setup_connections!
       end
     end
 
-    def self.with_connections
+    def self.with_servers
       servers = setup_connections!
 
       if !GlobalUid.configuration.per_process_affinity?
@@ -79,7 +79,7 @@ def self.with_connections
       errors = []
       servers.each do |server|
         begin
-          yield server.connection unless server.connection.nil?
+          yield server unless server.connection.nil?
         rescue TimeoutException, Exception => e
           GlobalUid.configuration.notifier.notify(e, e.message)
           errors << e
@@ -89,12 +89,12 @@ def self.with_connections
       end
 
       # in the case where all servers are gone, put everyone back in.
-      if servers.all? { |server| server.connection.nil? }
+      if get_connections.empty?
         servers.each { |server| server.update_retry_at(0) }
         raise NoServersAvailableException, "Errors hit: #{errors.map(&:to_s).join(',')}"
       end
 
-      servers.map { |server| server.connection }.compact
+      servers
     end
 
     def self.get_connections
@@ -103,8 +103,7 @@ def self.get_connections
     end
 
     def self.get_uid_for_class(klass)
-      with_connections do |connection|
-        server = self.servers.find { |s| connection.current_database.include?(s.name) }
+      with_servers do |server|
         Timeout.timeout(GlobalUid.configuration.query_timeout, TimeoutException) do
           return server.allocator.add(klass.global_uid_table)
         end
@@ -114,8 +113,7 @@ def self.get_uid_for_class(klass)
 
     def self.get_many_uids_for_class(klass, count)
       return [] unless count > 0
-      with_connections do |connection|
-        server = self.servers.find { |s| connection.current_database.include?(s.name) }
+      with_servers do |server|
         Timeout.timeout(GlobalUid.configuration.query_timeout, TimeoutException) do
           return server.allocator.add_many(klass.global_uid_table, count: count)
         end
diff --git a/test/lib/global_uid_test.rb b/test/lib/global_uid_test.rb
index b9c213a..cc5eda4 100644
--- a/test/lib/global_uid_test.rb
+++ b/test/lib/global_uid_test.rb
@@ -64,21 +64,21 @@ def table_exists?(connection, table)
         end
 
         it "create the global_uids table" do
-          GlobalUid::Base.with_connections do |cx|
-            assert table_exists?(cx, 'with_global_uids_ids'), 'Table should exist'
+          GlobalUid::Base.with_servers do |server|
+            assert table_exists?(server.connection, 'with_global_uids_ids'), 'Table should exist'
           end
         end
 
         it "create global_uids tables with matching ids" do
-          GlobalUid::Base.with_connections do |cx|
-            foo = cx.select_all("select id from with_global_uids_ids")
+          GlobalUid::Base.with_servers do |server|
+            foo = server.connection.select_all("select id from with_global_uids_ids")
             assert_equal(foo.first['id'].to_i, 1)
           end
         end
 
         it "create tables with the given storage_engine" do
-          GlobalUid::Base.with_connections do |cx|
-            foo = cx.select_all("show create table with_global_uids_ids")
+          GlobalUid::Base.with_servers do |server|
+            foo = server.connection.select_all("show create table with_global_uids_ids")
             assert_match(/ENGINE=InnoDB/, foo.first.values.join)
           end
         end
@@ -100,13 +100,13 @@ def table_exists?(connection, table)
       describe "dropping a table" do
         it "not drop the global-uid tables" do
           CreateWithNoParams.up
-          GlobalUid::Base.with_connections do |cx|
-            assert table_exists?(cx, 'with_global_uids_ids'), 'Table should exist'
+          GlobalUid::Base.with_servers do |server|
+            assert table_exists?(server.connection, 'with_global_uids_ids'), 'Table should exist'
           end
 
           CreateWithNoParams.down
-          GlobalUid::Base.with_connections do |cx|
-            assert table_exists?(cx, 'with_global_uids_ids'), 'Table should be dropped'
+          GlobalUid::Base.with_servers do |server|
+            assert table_exists?(server.connection, 'with_global_uids_ids'), 'Table should be dropped'
           end
         end
       end
@@ -118,8 +118,8 @@ def table_exists?(connection, table)
         end
 
         it "not create the global_uids table" do
-          GlobalUid::Base.with_connections do |cx|
-            assert !table_exists?(cx, 'with_global_uids_ids'), 'Table should not have been created'
+          GlobalUid::Base.with_servers do |server|
+            assert !table_exists?(server.connection, 'with_global_uids_ids'), 'Table should not have been created'
           end
         end
 
@@ -133,13 +133,13 @@ def table_exists?(connection, table)
       describe "dropping a table" do
         it "drop the global-uid tables" do
           CreateWithExplicitUidTrue.up
-          GlobalUid::Base.with_connections do |cx|
-            assert table_exists?(cx, 'with_global_uids_ids'), 'Table should exist'
+          GlobalUid::Base.with_servers do |server|
+            assert table_exists?(server.connection, 'with_global_uids_ids'), 'Table should exist'
           end
 
           CreateWithExplicitUidTrue.down
-          GlobalUid::Base.with_connections do |cx|
-            assert !table_exists?(cx, 'with_global_uids_ids'), 'Table should be dropped'
+          GlobalUid::Base.with_servers do |server|
+            assert !table_exists?(server.connection, 'with_global_uids_ids'), 'Table should be dropped'
           end
         end
       end
@@ -152,8 +152,8 @@ def table_exists?(connection, table)
       end
 
       it "not create the global_uids table" do
-        GlobalUid::Base.with_connections do |cx|
-          assert !table_exists?(cx, 'without_global_uids_ids'), 'Table should not not have been created'
+        GlobalUid::Base.with_servers do |server|
+          assert !table_exists?(server.connection, 'without_global_uids_ids'), 'Table should not not have been created'
         end
       end
 
@@ -289,9 +289,9 @@ def table_exists?(connection, table)
 
     describe "normally" do
       it "create tables with the default MyISAM storage engine" do
-        GlobalUid::Base.with_connections do |cx|
-          foo = cx.select_all("show create table with_global_uids_ids")
-          assert_match(/ENGINE=MyISAM/, foo.first.values.join)
+        GlobalUid::Base.with_servers do |server|
+          table = server.connection.select_all("show create table with_global_uids_ids")
+          assert_match(/ENGINE=MyISAM/, table.first.values.join)
         end
       end
 
@@ -315,8 +315,8 @@ def table_exists?(connection, table)
           end
 
           it "raises an exception, preventing duplicate ID generation" do
-            GlobalUid::Base.with_connections do |con|
-              con.execute("SET SESSION auto_increment_increment = 42")
+            GlobalUid::Base.with_servers do |server|
+              server.connection.execute("SET SESSION auto_increment_increment = 42")
             end
 
             assert_raises(GlobalUid::NoServersAvailableException) { test_unique_ids(10) }
@@ -324,8 +324,8 @@ def table_exists?(connection, table)
           end
 
           it "raises an exception before attempting to generate many UIDs" do
-            GlobalUid::Base.with_connections do |con|
-              con.execute("SET SESSION auto_increment_increment = 42")
+            GlobalUid::Base.with_servers do |server|
+              server.connection.execute("SET SESSION auto_increment_increment = 42")
             end
 
             assert_raises GlobalUid::NoServersAvailableException do
@@ -335,8 +335,8 @@ def table_exists?(connection, table)
           end
 
           it "doesn't cater for increment_by being increased by a factor of x" do
-            GlobalUid::Base.with_connections do |connection|
-              connection.execute("SET SESSION auto_increment_increment = #{GlobalUid.configuration.increment_by * 2}")
+            GlobalUid::Base.with_servers do |server|
+              server.connection.execute("SET SESSION auto_increment_increment = #{GlobalUid.configuration.increment_by * 2}")
             end
             # Due to multiple processes and threads sharing the same alloc server, identifiers may be provisioned
             # before the current thread receives its next one. We rely on the gap being divisible by the configured increment
@@ -448,8 +448,8 @@ def with_timed_out_connection(server:, duration:)
 
     describe "With both servers throwing exceptions" do
       before do
-        GlobalUid::Base.with_connections do |cx|
-          cx.stubs(:insert).raises(ActiveRecord::StatementInvalid)
+        GlobalUid::Base.with_servers do |server|
+          server.connection.stubs(:insert).raises(ActiveRecord::StatementInvalid)
         end
       end