mysql-proxy 读写分离脚本 ytjwt

local commands    = require("proxy.commands")
local tokenizer   = require("proxy.tokenizer")
local lb          = require("proxy.balance")
--allow host IP
--local host_allow={"10.50.21.100","127.0.0.1"}
if not proxy.global.config.rwsplit then
proxy.global.config.rwsplit = {
   max_idle_connections = 30,
   is_debug = true
}
end

local is_prepare                =false
local is_set                    =false

function connect_server()
local is_debug = proxy.global.config.rwsplit.is_debug
local rw_ndx = 1
local min_cur_idle_connections=1000000
      -- init all backends
      for i = 1, #proxy.global.backends do
         local s        = proxy.global.backends[i]
         local pool     = s.pool -- we don't have a username yet, try to find a connections which is idling
         local cur_idle = pool.users[""].cur_idle_connections
         if min_cur_idle_connections > cur_idle
         then
         min_cur_idle_connections=cur_idle
         rw_ndx=i
         end
       end

      proxy.connection.backend_ndx = rw_ndx
if is_debug then
   print("[connect_server] " .. proxy.connection.client.src.name)
   print("[server] " .. proxy.global.backends[proxy.connection.backend_ndx].dst.name)
end
end
function read_auth()
local is_debug = proxy.global.config.rwsplit.is_debug
local allow_host=true
--for i,host_allow_ip in ipairs(host_allow)
--do
--   if proxy.connection.client.src.address == host_allow_ip
--     then
--        allow_host=true
--         break
--     end
--end
if debug
   then
    print("[read_auth] client IP: " .. proxy.connection.client.src.address)
end      
   if not allow_host
   then
      proxy.response.type = proxy.MYSQLD_PACKET_ERR
                proxy.response.errmsg = "Logins are not allowed"
                  return proxy.PROXY_SEND_RESULT
   end
end
function read_auth_result( auth )
local is_debug = proxy.global.config.rwsplit.is_debug
if is_debug then
   print("[read_auth_result] " .. proxy.connection.client.src.name)
end
if auth.packet:byte() == proxy.MYSQLD_PACKET_OK then
   -- auth was fine, disconnect from the server
   proxy.connection.backend_ndx = 0
elseif auth.packet:byte() == proxy.MYSQLD_PACKET_EOF then
   -- we received either a
   --
   -- * MYSQLD_PACKET_ERR and the auth failed or
   -- * MYSQLD_PACKET_EOF which means a OLD PASSWORD (4.0) was sent
   print("(read_auth_result) ... not ok yet");
elseif auth.packet:byte() == proxy.MYSQLD_PACKET_ERR then
   -- auth failed
end
end

function read_query( packet )
local is_debug = proxy.global.config.rwsplit.is_debug
local cmd      = commands.parse(packet)
local c        = proxy.connection.client

local tokens

is_prepare = false  
is_set = false

-- looks like we have to forward this statement to a backend
if is_debug then
   print("[read_query] " .. proxy.connection.client.src.name)
   print(" current backend   = " .. proxy.connection.backend_ndx)
   print(" client default db = " .. c.default_db)
   print(" client username   = " .. c.username)
   if cmd.type == proxy.COM_QUERY then
    print(" query             = "        .. cmd.query)
   end
end

   if cmd.type == proxy.COM_QUIT then
   -- don't send COM_QUIT to the backend. We manage the connection
   -- in all aspects.
   proxy.response = {
    type = proxy.MYSQLD_PACKET_OK,
   }

   if is_debug then
    print(" (QUIT) current backend   = " .. proxy.connection.backend_ndx)
   end
   return proxy.PROXY_SEND_RESULT
end

 -- COM_BINLOG_DUMP packet can't be balanced
 --
 -- so we must send it always to the master
 if cmd.type == proxy.COM_BINLOG_DUMP then
  -- if we don't have a backend selected, let's pick the master
  --
  if proxy.connection.backend_ndx == 0 then
   proxy.connection.backend_ndx = lb.idle_failsafe_rw()
  end

  return
 end
 
 proxy.queries:append(1, packet, { resultset_is_needed = true })

 -- read/write splitting
 --
 -- send all non-transactional SELECTs to a slave
if not is_in_transaction and
    cmd.type == proxy.COM_QUERY then
   tokens     = tokens or assert(tokenizer.tokenize(cmd.query))
   local stmt = tokenizer.first_stmt_token(tokens)
   
   
    if stmt.token_name == "TK_SQL_SELECT" then
   
       local is_insert_id = false
     
       for i = 1, #tokens do
        local token = tokens[i]
          local utext = token.text:upper()
          if token.token_name == "TK_LITERAL"
             then
              if utext == "@@LAST_INSERT_ID" then
               is_insert_id = true
              end
         elseif token.token_name == "TK_FUNCTION"
          then
           if utext == "LAST_INSERT_ID" or utext == "IS_FREE_LOCK" or utext == "IS_USED_LOCK" or utext == "RELEASE_LOCK" or utext == "GET_LOCK"
             then
              is_insert_id = true
              end
          end     
   
       end
         
         
    if not is_insert_id     --如果是查询
    then local backend_ndx = lb.idle_ro()
    
        if backend_ndx > 0 then
         proxy.connection.backend_ndx = backend_ndx
        end
       else
        print("   found a SELECT LAST_INSERT_ID(), staying on the same backend")
       end
    
   
      elseif not is_prepare and stmt.token_name == "TK_LITERAL"
        then
           local token = tokens[1]
           local utext = token.text:upper()
             if utext == "PREPARE" or utext=="DEALLOCATE" or utext=="EXECUTE "
             then         
               is_prepare=true
             end  
             
      elseif not is_set and stmt.token_name == "TK_SQL_SET"
         then            
           is_set=true          
       end
   elseif cmd.type ==proxy.COM_STMT_PREPARE or cmd.type==proxy.COM_STMT_EXECUTE or cmd.type==proxy.COM_STMT_CLOSE
      then
       is_prepare=true                                                
end

   if proxy.connection.backend_ndx == 0 then
   -- we don't have a backend right now
   --
   -- let's pick a master as a good default
   --
   proxy.connection.backend_ndx = lb.idle_failsafe_rw()
end

-- by now we should have a backend
--
-- in case the master is down, we have to close the client connections
-- otherwise we can go on
if proxy.connection.backend_ndx == 0 then
   return proxy.PROXY_SEND_QUERY
end
local s = proxy.connection.server
 
   if cmd.type ~= proxy.COM_INIT_DB and
    c.default_db and c.default_db ~= s.default_db
     then
             print("    server default db: " .. s.default_db)
             print("    client default db: " .. c.default_db)
             print("    syncronizing")
             proxy.queries:append(2, string.char(proxy.COM_INIT_DB) .. c.default_db, { resultset_is_needed = true })
              proxy.queries:append(1, packet, { resultset_is_needed = true })
      else             
            
             print("    server default db: " .. s.default_db)
            print("    client default db: " .. c.default_db)
            proxy.queries:append(1, packet, { resultset_is_needed = true })
end

   -- send to master
if is_debug then
   if proxy.connection.backend_ndx > 0 then
    local b = proxy.global.backends[proxy.connection.backend_ndx]
    print("    sending to backend : " .. b.dst.name);
    print("    is_read_only[slave server]         : " .. tostring(b.type == proxy.BACKEND_TYPE_RO));
    print("    server default db: " .. s.default_db)
    print("    server username : " .. s.username)  
    print("    server thread_id :"..s.thread_id)
   end
   print("    in_trans        : " .. tostring(is_in_transaction))
   print("    COM_QUERY       : " .. tostring(cmd.type == proxy.COM_QUERY))
end
return proxy.PROXY_SEND_QUERY
end

function read_query_result( inj )
local is_debug = proxy.global.config.rwsplit.is_debug
local res      = assert(inj.resultset)
local flags    = res.flags
if inj.id ~= 1 then
   -- ignore the result of the USE <default_db>
   -- the DB might not exist on the backend, what do do ?
   --
   if inj.id == 2 then
    -- the injected INIT_DB failed as the slave doesn't have this DB
    -- or doesn't have permissions to read from it
    if res.query_status == proxy.MYSQLD_PACKET_ERR then
     proxy.queries:reset()
     proxy.response = {
      type = proxy.MYSQLD_PACKET_ERR,
      errmsg = "can't change DB ".. proxy.connection.client.default_db ..
       " to on slave " .. proxy.global.backends[proxy.connection.backend_ndx].dst.name
     }
     return proxy.PROXY_SEND_RESULT
    end
   end
   return proxy.PROXY_IGNORE_RESULT
end
is_in_transaction = flags.in_trans
local have_last_insert_id = (res.insert_id and (res.insert_id > 0))
if not is_in_transaction and
    not have_last_insert_id and
    not is_prepare and
    not is_set then
   -- release the backend
   proxy.connection.backend_ndx = 0
elseif is_debug then
   print("(read_query_result) staying on the same backend")
   print("    in_trans        : " .. tostring(is_in_transaction))
   print("    have_insert_id : " .. tostring(have_last_insert_id))
   print("have prepare :"..tostring(is_prepare))
   print("have set wrong :"..tostring(is_set))
end                           
end

function disconnect_client()
local is_debug = proxy.global.config.rwsplit.is_debug
if is_debug then
   print("[disconnect_client] " .. proxy.connection.client.src.name)
end
-- make sure we are disconnection from the connection
-- to move the connection into the pool
    if proxy.connection.backend_ndx == 0
    then
       for i = 1, #proxy.global.backends do
       local s = proxy.global.backends[i]
          local pool     = s.pool
         local cur_idle = pool.users[""].cur_idle_connections
         pool.max_idle_connections = proxy.global.config.rwsplit.max_idle_connections
         if is_debug then
            print ("cur_idle="..cur_idle )
            print ("pool.max_idle_connections = "..pool.max_idle_connections)
        end
         if s.state ~= proxy.BACKEND_STATE_DOWN and
            cur_idle > pool.max_idle_connections then
          -- try to disconnect a backend
           proxy.connection.backend_ndx = 0
           return
          end
        end
      end 
end