activerecord_partitioning 0.0.1 → 0.0.2

data/lib/activerecord_partitioning/connection_pools.rb

@@ -3,19 +3,26 @@ module ActiveRecordPartitioning
   end
 
   class ConnectionPools
-    attr_reader :store
+    attr_reader :key_name, :store
 
-    def initialize(store={})
+    def initialize(key_name, store={})
+      @key_name = key_name
       @store = store
     end
 
-    def [](key)
+    def merge!(pools)
+      pools.each do |klass_name, pool|
+        self[klass_name] = pool
+      end
+    end
+
+    def [](klass_name)
       config = ActiveRecordPartitioning.current_connection_pool_config
       raise NoActiveConnectionPoolError if config.nil?
       @store[connection_pool_key(config)]
     end
 
-    def []=(key, pool)
+    def []=(klass_name, pool)
       @store[connection_pool_key(pool.spec.config)] = pool
     end
 
@@ -33,7 +40,7 @@ module ActiveRecordPartitioning
 
     private
     def connection_pool_key(config)
-      config[:url]
+      config[@key_name]
     end
   end
 end

data/lib/activerecord_partitioning.rb

@@ -3,9 +3,15 @@ require 'activerecord_partitioning/connection_pools'
 
 module ActiveRecordPartitioning
   module_function
-  def setup(base_config = {}, pools = {})
+  def setup(key_name, base_config = {}, store = {})
     @base_config = base_config.symbolize_keys
-    ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new(ConnectionPools.new(pools))
+    new_pools = ConnectionPools.new(key_name, store)
+    new_pools.merge!(ActiveRecord::Base.connection_handler.connection_pools)
+    ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new(new_pools)
+  end
+
+  def reset_connection_handler
+    ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
   end
 
   def with_connection_pool(config, &block)

metadata

@@ -1,13 +1,13 @@
 --- !ruby/object:Gem::Specification 
 name: activerecord_partitioning
 version: !ruby/object:Gem::Version 
-  hash: 29
+  hash: 27
   prerelease: 
   segments: 
   - 0
   - 0
-  - 1
-  version: 0.0.1
+  - 2
+  version: 0.0.2
 platform: ruby
 authors: 
 - Xiao Li