viking-sequel 3.10.0

data/lib/sequel/adapters/utils/stored_procedures.rb

@@ -0,0 +1,84 @@
+module Sequel 
+  class Dataset
+    module StoredProcedureMethods
+      SQL_QUERY_TYPE = Hash.new{|h,k| h[k] = k}
+      SQL_QUERY_TYPE[:first] = SQL_QUERY_TYPE[:all] = :select
+      
+      # The name of the stored procedure to call
+      attr_accessor :sproc_name
+      
+      # The name of the stored procedure to call
+      attr_writer :sproc_args
+      
+      # Call the stored procedure with the given args
+      def call(*args, &block)
+        sp = clone
+        sp.sproc_args = args
+        sp.run(&block)
+      end
+
+      # Programmer friendly string showing this is a stored procedure,
+      # showing the name of the procedure.
+      def inspect
+        "<#{self.class.name}/StoredProcedure name=#{@sproc_name}>"
+      end
+      
+      # Run the stored procedure with the current args on the database
+      def run(&block)
+        case @sproc_type
+        when :select, :all
+          all(&block)
+        when :first
+          first
+        when :insert
+          insert
+        when :update
+          update
+        when :delete
+          delete
+        end
+      end
+      
+      # Set the type of the stored procedure and override the corresponding _sql
+      # method to return the empty string (since the result will be
+      # ignored anyway).
+      def sproc_type=(type)
+        @sproc_type = type
+        meta_def("#{sql_query_type}_sql"){|*a| ''}
+      end
+      
+      private
+      
+      # The type of query (:select, :insert, :delete, :update).
+      def sql_query_type
+        SQL_QUERY_TYPE[@sproc_type]
+      end
+    end
+  
+    module StoredProcedures
+      # For the given type (:select, :first, :insert, :update, or :delete),
+      # run the database stored procedure with the given name with the given
+      # arguments.
+      def call_sproc(type, name, *args)
+        prepare_sproc(type, name).call(*args)
+      end
+      
+      # Transform this dataset into a stored procedure that you can call
+      # multiple times with new arguments.
+      def prepare_sproc(type, name)
+        sp = clone
+        prepare_extend_sproc(sp)
+        sp.sproc_type = type
+        sp.sproc_name = name
+        sp
+      end
+      
+      private
+      
+      # Extend the dataset with the stored procedure methods.
+      def prepare_extend_sproc(ds)
+        ds.extend(StoredProcedureMethods)
+      end
+    end
+  end
+end

data/lib/sequel/connection_pool.rb

@@ -0,0 +1,99 @@
+# The base connection pool class, which all other connection pools are built
+# on.  This class is not instantiated directly, but subclasses should at
+# the very least implement the following API:
+# * initialize(Hash, &block) - The block is used as the connection proc,
+#   which should accept a single symbol argument.
+# * hold(Symbol, &block) - yield a connection object (obtained from calling
+#   the block passed to initialize) to the current block. For sharded
+#   connection pools, the Symbol passed is the shard/server to use.
+# * disconnect(Symbol, &block) - disconnect the connection object.  If a
+#   block is given, pass the connection option to it, otherwise use the
+#   :disconnection_proc option in the hash passed to initialize.  For sharded
+#   connection pools, the Symbol passed is the shard/server to use.
+# * servers - an array of shard/server symbols for all shards/servers that this
+#   connection pool recognizes.
+# * size - an integer representing the total number of connections in the pool,
+#   or for the given shard/server if sharding is supported.
+#
+# For sharded connection pools, the sharded API:
+# * add_servers(Array of Symbols) - start recognizing all shards/servers specified
+#   by the array of symbols.
+# * remove_servers(Array of Symbols) - no longer recognize all shards/servers
+#   specified by the array of symbols.
+class Sequel::ConnectionPool
+  # The default server to use
+  DEFAULT_SERVER = :default
+  
+  # A map of [single threaded, sharded] values to files (indicating strings to
+  # be required) ConnectionPool subclasses.
+  CONNECTION_POOL_MAP = {[true, false] => :single, 
+    [true, true] => :sharded_single,
+    [false, false] => :threaded,
+    [false, true] => :sharded_threaded}
+  
+  # Class methods used to return an appropriate pool subclass, separated
+  # into a module for easier overridding by extensions.
+  module ClassMethods
+    # Return a pool subclass instance based on the given options.  If a :pool_class
+    # option is provided is provided, use that pool class, otherwise
+    # use a new instance of an appropriate pool subclass based on the
+    # :single_threaded and :servers options.
+    def get_pool(opts = {}, &block)
+      case v = connection_pool_class(opts)
+      when Class
+        v.new(opts, &block)
+      when Symbol
+        Sequel.ts_require("connection_pool/#{v}")
+        connection_pool_class(opts).new(opts, &block) || raise(Sequel::Error, "No connection pool class found")
+      end
+    end
+    
+    private
+    
+    # Return a connection pool class based on the given options.
+    def connection_pool_class(opts)
+      opts[:pool_class] || CONNECTION_POOL_MAP[[!!opts[:single_threaded], !!opts[:servers]]]
+    end
+  end
+  extend ClassMethods
+  
+  # Instantiates a connection pool with the given options.  The block is called
+  # with a single symbol (specifying the server/shard to use) every time a new
+  # connection is needed.  The following options are respected for all connection
+  # pools:
+  # * :after_connect - The proc called after each new connection is made, with the
+  #   connection object, useful for customizations that you want to apply to all
+  #   connections.
+  # * :disconnection_proc - The proc called when removing connections from the pool,
+  #   which is passed the connection to disconnect.
+  def initialize(opts={}, &block)
+    raise(Sequel::Error, "No connection proc specified") unless @connection_proc = block
+    @disconnection_proc = opts[:disconnection_proc]
+    @after_connect = opts[:after_connect]
+  end
+  
+  # Alias for size, not aliased directly for ease of subclass implementation
+  def created_count(*args)
+    size(*args)
+  end
+  
+  # An array of symbols for all shards/servers, which is a single :default by default.
+  def servers
+    [:default]
+  end
+  
+  private
+  
+  # Return a new connection by calling the connection proc with the given server name,
+  # and checking for connection errors.
+  def make_new(server)
+    begin
+      conn = @connection_proc.call(server)
+      @after_connect.call(conn) if @after_connect
+    rescue Exception=>exception
+      raise Sequel.convert_exception_class(exception, Sequel::DatabaseConnectionError)
+    end
+    raise(Sequel::DatabaseConnectionError, "Connection parameters not valid") unless conn
+    conn
+  end
+end

data/lib/sequel/connection_pool/sharded_single.rb

@@ -0,0 +1,84 @@
+# A ShardedSingleConnectionPool is a single threaded connection pool that
+# works with multiple shards/servers.
+class Sequel::ShardedSingleConnectionPool < Sequel::ConnectionPool
+  # Initializes the instance with the supplied block as the connection_proc.
+  #
+  # The single threaded pool takes the following options:
+  #
+  # * :servers - A hash of servers to use.  Keys should be symbols.  If not
+  #   present, will use a single :default server.  The server name symbol will
+  #   be passed to the connection_proc.
+  def initialize(opts={}, &block)
+    super
+    @conns = {}
+    @servers = Hash.new(:default)
+    add_servers([:default])
+    add_servers(opts[:servers].keys) if opts[:servers]
+  end
+  
+  # Adds new servers to the connection pool. Primarily used in conjunction with master/slave
+  # or shard configurations. Allows for dynamic expansion of the potential slaves/shards
+  # at runtime. servers argument should be an array of symbols. 
+  def add_servers(servers)
+    servers.each{|s| @servers[s] = s}
+  end
+  
+  # The connection for the given server.
+  def conn(server=:default)
+    @conns[@servers[server]]
+  end
+  
+  # Disconnects from the database. Once a connection is requested using
+  # #hold, the connection is reestablished. Options:
+  # * :server - Should be a symbol specifing the server to disconnect from,
+  #   or an array of symbols to specify multiple servers.
+  def disconnect(opts={}, &block)
+    block ||= @disconnection_proc
+    (opts[:server] ? Array(opts[:server]) : servers).each{|s| disconnect_server(s, &block)}
+  end
+  
+  # Yields the connection to the supplied block for the given server.
+  # This method simulates the ConnectionPool#hold API.
+  def hold(server=:default)
+    begin
+      server = @servers[server]
+      yield(@conns[server] ||= make_new(server))
+    rescue Sequel::DatabaseDisconnectError
+      disconnect_server(server, &@disconnection_proc)
+      raise
+    end
+  end
+  
+  # Remove servers from the connection pool. Primarily used in conjunction with master/slave
+  # or shard configurations.  Similar to disconnecting from all given servers,
+  # except that after it is used, future requests for the server will use the
+  # :default server instead.
+  def remove_servers(servers)
+    raise(Sequel::Error, "cannot remove default server") if servers.include?(:default)
+    servers.each do |server|
+      disconnect_server(server, &@disconnection_proc)
+      @servers.delete(server)
+    end
+  end
+  
+  # Return an array of symbols for servers in the connection pool.
+  def servers
+    @servers.keys
+  end
+  
+  # The number of different shards/servers this pool is connected to.
+  def size
+    @conns.length
+  end
+  
+  private
+  
+  # Disconnect from the given server, if connected.
+  def disconnect_server(server, &block)
+    if conn = @conns.delete(server)
+      block.call(conn) if block
+    end
+  end
+  
+  CONNECTION_POOL_MAP[[true, true]] = self
+end

data/lib/sequel/connection_pool/sharded_threaded.rb

@@ -0,0 +1,211 @@
+Sequel.require 'connection_pool/threaded'
+
+# The slowest and most advanced connection, dealing with both multi-threaded
+# access and configurations with multiple shards/servers.
+#
+# In addition, this pool subclass also handles scheduling in-use connections
+# to be removed from the pool when they are returned to it.
+class Sequel::ShardedThreadedConnectionPool < Sequel::ThreadedConnectionPool
+  # The following additional options are respected:
+  # * :servers - A hash of servers to use.  Keys should be symbols.  If not
+  #   present, will use a single :default server.  The server name symbol will
+  #   be passed to the connection_proc.
+  def initialize(opts = {}, &block)
+    super
+    @available_connections = {}
+    @connections_to_remove = []
+    @servers = Hash.new(:default)
+    add_servers([:default])
+    add_servers(opts[:servers].keys) if opts[:servers]
+  end
+  
+  # Adds new servers to the connection pool. Primarily used in conjunction with master/slave
+  # or shard configurations. Allows for dynamic expansion of the potential slaves/shards
+  # at runtime. servers argument should be an array of symbols. 
+  def add_servers(servers)
+    sync do
+      servers.each do |server|
+        unless @servers.has_key?(server)
+          @servers[server] = server
+          @available_connections[server] = []
+          @allocated[server] = {}
+        end
+      end
+    end
+  end
+  
+  # A hash of connections currently being used for the given server, key is the
+  # Thread, value is the connection.  Nonexistent servers will return nil.  Treat
+  # this as read only, do not modify the resulting object.
+  def allocated(server=:default)
+    @allocated[server]
+  end
+  
+  # An array of connections opened but not currently used, for the given
+  # server. Nonexistent servers will return nil. Treat this as read only, do
+  # not modify the resulting object.
+  def available_connections(server=:default)
+    @available_connections[server]
+  end
+  
+  # The total number of connections opened for the given server, should
+  # be equal to available_connections.length + allocated.length.  Nonexistent
+  # servers will return the created count of the default server.
+  def size(server=:default)
+    server = @servers[server]
+    @allocated[server].length + @available_connections[server].length
+  end
+  
+  # Removes all connection currently available on all servers, optionally
+  # yielding each connection to the given block. This method has the effect of 
+  # disconnecting from the database, assuming that no connections are currently
+  # being used.  If connections are being used, they are scheduled to be
+  # disconnected as soon as they are returned to the pool.
+  # 
+  # Once a connection is requested using #hold, the connection pool
+  # creates new connections to the database. Options:
+  # * :server - Should be a symbol specifing the server to disconnect from,
+  #   or an array of symbols to specify multiple servers.
+  def disconnect(opts={}, &block)
+    block ||= @disconnection_proc
+    sync do
+      (opts[:server] ? Array(opts[:server]) : @servers.keys).each do |s|
+        disconnect_server(s, &block)
+      end
+    end
+  end
+  
+  # Chooses the first available connection to the given server, or if none are
+  # available, creates a new connection.  Passes the connection to the supplied
+  # block:
+  # 
+  #   pool.hold {|conn| conn.execute('DROP TABLE posts')}
+  # 
+  # Pool#hold is re-entrant, meaning it can be called recursively in
+  # the same thread without blocking.
+  #
+  # If no connection is immediately available and the pool is already using the maximum
+  # number of connections, Pool#hold will block until a connection
+  # is available or the timeout expires.  If the timeout expires before a
+  # connection can be acquired, a Sequel::PoolTimeout is 
+  # raised.
+  def hold(server=:default)
+    sync{server = @servers[server]}
+    t = Thread.current
+    if conn = owned_connection(t, server)
+      return yield(conn)
+    end
+    begin
+      unless conn = acquire(t, server)
+        time = Time.now
+        timeout = time + @timeout
+        sleep_time = @sleep_time
+        sleep sleep_time
+        until conn = acquire(t, server)
+          raise(::Sequel::PoolTimeout) if Time.now > timeout
+          sleep sleep_time
+        end
+      end
+      yield conn
+    rescue Sequel::DatabaseDisconnectError
+      sync{@connections_to_remove << conn} if conn
+      raise
+    ensure
+      sync{release(t, conn, server)} if conn
+    end
+  end
+
+  # Remove servers from the connection pool. Primarily used in conjunction with master/slave
+  # or shard configurations.  Similar to disconnecting from all given servers,
+  # except that after it is used, future requests for the server will use the
+  # :default server instead.
+  def remove_servers(servers)
+    sync do
+      raise(Sequel::Error, "cannot remove default server") if servers.include?(:default)
+      servers.each do |server|
+        if @servers.include?(server)
+          disconnect_server(server, &@disconnection_proc)
+          @available_connections.delete(server)
+          @allocated.delete(server)
+          @servers.delete(server)
+        end
+      end
+    end
+  end
+
+  # Return an array of symbols for servers in the connection pool.
+  def servers
+    sync{@servers.keys}
+  end
+
+  private
+
+  # Assigns a connection to the supplied thread for the given server, if one
+  # is available. The calling code should NOT already have the mutex when
+  # calling this.
+  def acquire(thread, server)
+    sync do
+      if conn = available(server)
+        allocated(server)[thread] = conn
+      end
+    end
+  end
+  
+  # Returns an available connection to the given server. If no connection is
+  # available, tries to create a new connection. The calling code should already
+  # have the mutex before calling this.
+  def available(server)
+    available_connections(server).pop || make_new(server)
+  end
+
+  # Disconnect from the given server.  Disconnects available connections
+  # immediately, and schedules currently allocated connections for disconnection
+  # as soon as they are returned to the pool. The calling code should already
+  # have the mutex before calling this.
+  def disconnect_server(server, &block)
+    if conns = available_connections(server)
+      conns.each{|conn| block.call(conn)} if block
+      conns.clear
+    end
+    @connections_to_remove.concat(allocated(server).values)
+  end
+
+  # Creates a new connection to the given server if the size of the pool for
+  # the server is less than the maximum size of the pool. The calling code
+  # should already have the mutex before calling this.
+  def make_new(server)
+    if (n = size(server)) >= @max_size
+      allocated(server).to_a.each{|t, c| release(t, c, server) unless t.alive?}
+      n = nil
+    end
+    default_make_new(server) if (n || size(server)) < @max_size
+  end
+  
+  # Returns the connection owned by the supplied thread for the given server,
+  # if any. The calling code should NOT already have the mutex before calling this.
+  def owned_connection(thread, server)
+    sync{@allocated[server][thread]}
+  end
+  
+  # Releases the connection assigned to the supplied thread and server. If the
+  # server or connection given is scheduled for disconnection, remove the
+  # connection instead of releasing it back to the pool.
+  # The calling code should already have the mutex before calling this.
+  def release(thread, conn, server)
+    if @connections_to_remove.include?(conn)
+      remove(thread, conn, server)
+    else
+      available_connections(server) << allocated(server).delete(thread)
+    end
+  end
+
+  # Removes the currently allocated connection from the connection pool. The
+  # calling code should already have the mutex before calling this.
+  def remove(thread, conn, server)
+    @connections_to_remove.delete(conn)
+    allocated(server).delete(thread) if @servers.include?(server)
+    @disconnection_proc.call(conn) if @disconnection_proc
+  end
+  
+  CONNECTION_POOL_MAP[[false, true]] = self
+end