dennybaa
8/31/2011 - 9:16 AM

map_dbnames

map_dbnames

-- Generated by Chef
---
-- Does simple SQF-type load balancing, ie servers with least connections
-- are used first. Provides mechanism for rewriting database names on fly.

---
-- config
---
-- database name mappings
local database_names = { "__OAG", "__IP2Location" }
local database_mapto = { "oag_201108", "IP2Location201108" }
---
local proto = assert(require("mysql.proto"))
local MAX_INT=65535

---
-- 
function connect_server()
  local least_conn = MAX_INT
  local least_conn_ndx = 0
  
  for i = 1, #proxy.global.backends do
    local s = proxy.global.backends[i]
    
    if s.state ~= proxy.BACKEND_STATE_DOWN and least_conn > s.connected_clients then
      least_conn = s.connected_clients
      least_conn_ndx = i
    end
  end
  proxy.connection.backend_ndx = least_conn_ndx
end

---
-- process query
function read_query( packet )
  local query = string.sub(packet, 2)
  local qtype = string.byte(packet)
  local c = proxy.connection.client
  ---
  -- rewrite db name on init
  if qtype == proxy.COM_INIT_DB then
    local mapped = map_dbname(string.gsub(query,"%s",""))
    if mapped ~= "" then
      proxy.queries:append(1, string.char(proxy.COM_INIT_DB) .. mapped )
      return proxy.PROXY_SEND_QUERY
    end
  ---
  -- rewrite default_db on query
  elseif qtype == proxy.COM_QUERY then
    local mapped = map_dbname(c.default_db)
    if mapped ~= "" then
      proxy.queries:append(2, string.char(proxy.COM_INIT_DB) .. mapped, { resultset_is_needed = true })
      proxy.queries:append(1, packet)
      return proxy.PROXY_SEND_QUERY
    end
  end
end

---
-- rewrite database name during authentification
function read_auth()
  local c = proxy.connection.client
  local mapped = map_dbname(c.default_db)
  if mapped ~= "" then
    local response_packet = proto.to_response_packet({
        username = c.username,
        response = c.scrambled_password,
        charset  = 8, -- default charset
        database = mapped,
        max_packet_size = 1 * 1024 * 1024
      })
    proxy.queries:append(1, response_packet)
    return proxy.PROXY_SEND_QUERY
  end
end

---
-- map virtual database name to a real one
function map_dbname(name)
  for i, v in ipairs(database_names) do
    if v == name then return database_mapto[i] end
  end
  return ""
end