Support connection load balancing

src/backend_message.coffee

@@ -10,6 +10,12 @@ class BackendMessage
   read: (buffer) ->
     # Implement me in subclass
 
+class BackendMessage.ConnectionLoadBalanceResponse extends BackendMessage
+  typeId: 89
+
+  read: (buffer) ->
+    this.host = buffer.toString('ascii', 4, buffer.length - 1)
+
 
 class BackendMessage.Authentication extends BackendMessage
   typeId: 82 # R

src/connection.coffee

@@ -28,14 +28,45 @@ class Connection extends EventEmitter
 
   connect: (callback) ->
     @connectedCallback = callback
-    @connection = net.createConnection @connectionOptions.port, @connectionOptions.host
 
     initialErrorHandler = (err) =>
       if @connectedCallback then @connectedCallback(err.message) else @emit 'error', err
 
-    @connection.on 'error', initialErrorHandler
+    connect = () =>
+      redirectConn = (msg) =>
+        console.log 'Redirect to = ', msg.host, @connectionOptions.port if @debug
+        @connection.destroy()
+        @redirctedHost = msg.host
+        @connection = net.createConnection(@connectionOptions.port, msg.host)
+        @connection.on 'error', initialErrorHandler
 
-    @connection.on 'connect', =>
+        @connection.on 'connect', () =>
+          @emit 'connect'
+
+      @connection = net.createConnection @connectionOptions.port, @connectionOptions.host
+
+      @connection.on 'error', initialErrorHandler
+
+      @connection.on 'connect', () =>
+        if @connectionOptions.loadBalance
+          @connection.removeListener 'error', initialErrorHandler
+          @_bindEventListeners()
+
+          @connection.once 'data', (buffer) =>
+            if buffer.toString('ascii') == 'N'
+              console.log "Connection load balancing is turned OFF on server" if @debug
+              @emit 'connect'
+            else
+              size = buffer.readUInt32BE(1)
+              message = BackendMessage.fromBuffer(buffer.slice(0, size + 1))
+              redirectConn message
+
+          @_writeMessage new FrontendMessage.ConnectionLoadBalance()
+
+        else
+          @emit 'connect'
+
+    @on 'connect', =>
       @connection.removeListener 'error', initialErrorHandler
       @connected = true
       @_bindEventListeners()
@@ -65,6 +96,10 @@ class Connection extends EventEmitter
       else
         @_handshake()
 
+    connect()
+
+    @connection
+
   _bindEventListeners: ->
     @connection.once 'close',   @_onClose.bind(this)
     @connection.once 'error',   @_onError.bind(this)

src/frontend_message.coffee

@@ -33,6 +33,13 @@ class FrontendMessage
 
     return messageBuffer
 
+class FrontendMessage.ConnectionLoadBalance extends FrontendMessage
+  typeId: null
+
+  payload: ->
+    arr = [0x0, 0x0, 0x0, 0x8, 0x4, 0xd3, 0x0, 0x0]
+    pl = new Buffer(arr)
+    pl.slice 4
 
 class FrontendMessage.Startup extends FrontendMessage
   typeId: null

test/functional/connection_test.coffee

@@ -71,6 +71,79 @@ describe 'Vertica.connect', ->
         assert.equal err.message, 'The connection was closed.'
         done()
 
+  describe 'Connection Load Balancing', ->
+    origLoadBalancingPolicy = ''
+
+    beforeEach (done) ->
+      if !fs.existsSync('./test/connection.json')
+        throw new Error("Create test/connection.json to run functional tests")
+      else
+        connectionInfo = JSON.parse(fs.readFileSync('./test/connection.json'))
+        Vertica.connect connectionInfo, (err, connection) ->
+          assert.equal err, null
+          assert.ok !connection.busy
+          assert.ok connection.connected
+
+          connection.query "SELECT GET_LOAD_BALANCE_POLICY()", (err, resultset) ->
+            assert.ok resultset instanceof Vertica.Resultset
+            assert.ok resultset.rows.length == 1
+            origLoadBalancingPolicy = resultset.rows[0][0]
+            done()
+
+    it "should connect even when load balancing is truned OFF", (done) ->
+      Vertica.connect connectionInfo, (err, connection) ->
+        assert.equal err, null
+        assert.ok !connection.busy
+        assert.ok connection.connected
+
+        connection.query "SELECT SET_LOAD_BALANCE_POLICY('NONE')", (err, resultset) ->
+          assert.equal err, null
+
+          connectionInfo.loadBalance = true
+          Vertica.connect connectionInfo, (err, conn) ->
+            assert.equal err, null
+            assert.ok !conn.busy
+            assert.ok conn.connected
+
+          done()
+
+        assert.ok connection.busy
+        assert.ok connection.connected
+
+    it "should connect to different host at least once", (done) ->
+      Vertica.connect connectionInfo, (err, connection) ->
+        assert.equal err, null
+        assert.ok !connection.busy
+        assert.ok connection.connected
+
+        connection.query "SELECT SET_LOAD_BALANCE_POLICY('ROUNDROBIN')", (err, resultset) ->
+          assert.equal err, null
+
+          connectionInfo.loadBalance = true
+          Vertica.connect connectionInfo, (err, conn) ->
+            assert.equal err, null
+            assert.ok !conn.busy
+            assert.ok conn.connected
+            if connectionInfo.host == conn.redirctedHost
+              Vertica.connect connectionInfo, (err, conn) ->
+                assert.equal err, null
+                assert.ok !conn.busy
+                assert.ok conn.connected
+                assert.notEqual conn.redirctedHost, connectionInfo.host
+
+          done()
+
+        assert.ok connection.busy
+        assert.ok connection.connected
+
+    afterEach (done) ->
+      Vertica.connect connectionInfo, (err, connection) ->
+        return done(err) if err?
+
+        connection.query "SELECT SET_LOAD_BALANCE_POLICY('" + origLoadBalancingPolicy + "')", (err, rs) ->
+          return done(err) if err?
+          done()
+
   describe 'Statement interruption', ->
     beforeEach (done) ->
       if !fs.existsSync('./test/connection.json')

test/unit/backend_message_test.coffee

@@ -2,6 +2,12 @@ assert         = require 'assert'
 BackendMessage = require('../../src/backend_message')
 Buffer         = require('../../src/buffer').Buffer
 
+describe 'BackendMessage.ConnectionLoadBalanceResponse', ->
+  it "should read a message correctly", ->
+    message = BackendMessage.fromBuffer(new Buffer([0x59, 0x00, 0x00, 0x00, 0x15, 0x00, 0x00, 0x15, 0x39, 0x31, 0x39, 0x32, 0x2e, 0x31, 0x36, 0x38, 0x2e, 0x31, 0x2e, 0x33, 0x39, 0x00]))
+    assert.ok message instanceof BackendMessage.ConnectionLoadBalanceResponse
+    assert.equal message.host, '192.168.1.39'
+
 describe 'BackendMessage.Authentication', ->
   it "should read a message correctly", ->
     message = BackendMessage.fromBuffer(new Buffer([0x52, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00]))

test/unit/frontend_message_test.coffee

@@ -4,6 +4,12 @@ Buffer          = require('../../src/buffer').Buffer
 FrontendMessage = require('../../src/frontend_message')
 Authentication  = require('../../src/authentication')
 
+describe "FrontendMessage.ConnectionLoadBalance", ->
+  it "should encode the message correctly", ->
+    topic = new FrontendMessage.ConnectionLoadBalance
+    reference = new Buffer([0x0, 0x0, 0x0, 0x8, 0x4, 0xD3, 0x0, 0x0])
+    assert.deepEqual topic.toBuffer(), reference
+
 describe "FrontendMessage.Startup", ->
   it "should hold the message's information", ->
     topic = new FrontendMessage.Startup('username', 'database')