binding.cc

#include <libpq-fe.h>
#include "type-oids.h"
#include <node.h>
#include <node_events.h>
#include <assert.h>
#include <stdlib.h>

using namespace v8;
using namespace node;
static Persistent<String> ready_symbol;
static Persistent<String> result_symbol;
static Persistent<String> close_symbol;
static Persistent<String> connect_symbol;
static Persistent<String> notify_symbol;
#define READY_STATE_SYMBOL String::NewSymbol("readyState")
#define MAP_TUPLE_ITEMS_SYMBOL String::NewSymbol("mapTupleItems")

class Connection : public EventEmitter {
 public:
  static void
  Initialize (v8::Handle<v8::Object> target)
  {
    HandleScope scope;

    Local<FunctionTemplate> t = FunctionTemplate::New(New);

    t->Inherit(EventEmitter::constructor_template);
    t->InstanceTemplate()->SetInternalFieldCount(1);

    close_symbol = NODE_PSYMBOL("close");
    connect_symbol = NODE_PSYMBOL("connect");
    result_symbol = NODE_PSYMBOL("result");
    ready_symbol = NODE_PSYMBOL("ready");
    notify_symbol = NODE_PSYMBOL("notify");

    NODE_SET_PROTOTYPE_METHOD(t, "connect", Connect);
    NODE_SET_PROTOTYPE_METHOD(t, "close", Close);
    NODE_SET_PROTOTYPE_METHOD(t, "reset", Reset);
    NODE_SET_PROTOTYPE_METHOD(t, "dispatchQuery", DispatchQuery);
    NODE_SET_PROTOTYPE_METHOD(t, "escapeString", EscapeString);

    t->PrototypeTemplate()->SetAccessor(READY_STATE_SYMBOL, ReadyStateGetter);
    t->PrototypeTemplate()->SetAccessor(MAP_TUPLE_ITEMS_SYMBOL,
                                        MapTupleItemsGetter,
                                        MapTupleItemsSetter);

    target->Set(String::NewSymbol("Connection"), t->GetFunction());
  }

  bool Connect (const char* conninfo)
  {
    if (connection_) return false;
    
    /* TODO Ensure that hostaddr is always specified in conninfo to avoid
     * name lookup. (Unless we're connecting to localhost.)
     */
    connection_ = PQconnectStart(conninfo);
    if (!connection_) return false;

    if (PQsetnonblocking(connection_, 1) == -1) {
      PQfinish(connection_);
      connection_ = NULL;
      return false;
    }

    ConnStatusType status = PQstatus(connection_);

    if (CONNECTION_BAD == status) {
      PQfinish(connection_);
      connection_ = NULL;
      return false;
    }

    connecting_ = true;

    int fd = PQsocket(connection_);

    ev_io_set(&read_watcher_, fd, EV_READ);
    ev_io_set(&write_watcher_, fd, EV_WRITE);

    /* If you have yet to call PQconnectPoll, i.e., just after the call to
     * PQconnectStart, behave as if it last returned PGRES_POLLING_WRITING.
     */
    ev_io_start(EV_DEFAULT_ &write_watcher_);
    
    Ref();

    return true;
  }

  bool Reset ( )
  {
    /* To initiate a connection reset, call PQresetStart. If it returns 0,
     * the reset has failed. If it returns 1, poll the reset using
     * PQresetPoll in exactly the same way as you would create the
     * connection using PQconnectPoll */
    int r = PQresetStart(connection_);
    if (r == 0) return false;

    if (PQsetnonblocking(connection_, 1) == -1) {
      PQfinish(connection_);
      connection_ = NULL;
      return false;
    }

    resetting_ = true;
    ev_io_set(&read_watcher_, PQsocket(connection_), EV_READ);
    ev_io_set(&write_watcher_, PQsocket(connection_), EV_WRITE);
    ev_io_start(EV_DEFAULT_ &write_watcher_);
    return true;
  }

  bool Query (const char *command) 
  {
    int r = PQsendQuery(connection_, command);
    if (r == 0) {
      return false;
    }
    if (PQflush(connection_) == 1) ev_io_start(EV_DEFAULT_ &write_watcher_);
    return true;
  }

  void Close (Local<Value> exception = Local<Value>())
  {
    HandleScope scope;
    ev_io_stop(EV_DEFAULT_ &write_watcher_);
    ev_io_stop(EV_DEFAULT_ &read_watcher_);
    PQfinish(connection_);
    connection_ = NULL;
    if (exception.IsEmpty()) {
      Emit(close_symbol, 0, NULL);
    } else {
      Emit(close_symbol, 1, &exception);
    }
    Unref();
  }

  char * ErrorMessage ( )
  {
    return PQerrorMessage(connection_);
  }

 protected:

  static Handle<Value>
  New (const Arguments& args)
  {
    HandleScope scope;

    Connection *connection = new Connection();
    connection->Wrap(args.This());

    return args.This();
  }

  static Handle<Value>
  EscapeString (const Arguments& args)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(args.This());

    HandleScope scope;

    if (args.Length() == 0 || !args[0]->IsString()) {
      return ThrowException(Exception::Error(
            String::New("Must give a string to escape")));
    }

    String::Utf8Value unescaped_string(args[0]->ToString());

    // the string to be escaped is a raw value and
    // may contain everything you can imagine
    int from_len = args[0]->ToString()->Utf8Length();

    char *to;
    to = (char *) malloc(from_len * 2 + 1);

    if (connection->connection_) {
      int error;
      // TODO handle an error if set
      PQescapeStringConn(connection->connection_,
                         to,
                         *unescaped_string,
                         (size_t)from_len,
                         &error);
    } else {
      PQescapeString(to, *unescaped_string, (size_t)from_len);
    }

    Handle<Value> result = String::New(to);

    free(to);

    return scope.Close(result);
  }

  static Handle<Value>
  Connect (const Arguments& args)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(args.This());

    HandleScope scope;

    if (args.Length() == 0 || !args[0]->IsString()) {
      return ThrowException(Exception::Error(
            String::New("Must give conninfo string as argument")));
    }

    String::Utf8Value conninfo(args[0]->ToString());

    bool r = connection->Connect(*conninfo);

    if (!r) {
      return ThrowException(Exception::Error(
            String::New(connection->ErrorMessage())));
    }

    return Undefined();
  }

  static Handle<Value>
  Close (const Arguments& args)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(args.This());
    HandleScope scope;
    connection->Close();
    return Undefined();
  }

  static Handle<Value>
  Reset (const Arguments& args)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(args.This());

    HandleScope scope;

    bool r = connection->Reset();

    if (!r) {
      return ThrowException(Exception::Error(
            String::New(connection->ErrorMessage())));
    }

    return Undefined();
  }

  static Handle<Value>
  DispatchQuery (const Arguments& args)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(args.This());

    HandleScope scope;

    if (args.Length() == 0 || !args[0]->IsString()) {
      return ThrowException(Exception::TypeError(
            String::New("First argument must be a string")));
    }

    String::Utf8Value query(args[0]->ToString());

    bool r = connection->Query(*query);

    if (!r) {
      return ThrowException(Exception::Error(
            String::New(connection->ErrorMessage())));
    }

    return Undefined();
  }

  static Handle<Value>
  MapTupleItemsGetter (Local<String> property, const AccessorInfo& info)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(info.This());
    assert(connection);

    assert(property == MAP_TUPLE_ITEMS_SYMBOL);

    HandleScope scope;

    return scope.Close(Boolean::New(connection->mapTupleItems_));
  }

  static void
  MapTupleItemsSetter (Local<String> property,
                       Local<Value> value,
                       const AccessorInfo& info)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(info.This());
    assert(connection);

    assert(property == MAP_TUPLE_ITEMS_SYMBOL);

    if (!value->IsBoolean()) {
      ThrowException(Exception::TypeError(
            String::New("mapTupleItems should be of Boolean value")));
    }

    connection->mapTupleItems_ = value->ToBoolean()->Value();
  }

  static Handle<Value>
  ReadyStateGetter (Local<String> property, const AccessorInfo& info)
  {
    Connection *connection = ObjectWrap::Unwrap<Connection>(info.This());
    assert(connection);

    assert(property == READY_STATE_SYMBOL);

    HandleScope scope;

    const char *s = NULL;

    switch(PQtransactionStatus(connection->connection_)) {
      case PQTRANS_IDLE:
        s = "OK";
        break;
      case PQTRANS_ACTIVE:
        s = "commandActive";
        break;
      case PQTRANS_UNKNOWN:
        s = "bad";
        break;
      case PQTRANS_INTRANS:
        s = "commandActiveButIdle";
        break;
    }

    if(s != NULL) {
        return scope.Close(String::NewSymbol(s));
    }

    switch (PQstatus(connection->connection_)) {
      case CONNECTION_STARTED: 
        s = "started";
        break;
      case CONNECTION_MADE: 
        s = "made";
        break;
      case CONNECTION_AWAITING_RESPONSE:
        s = "awaitingResponse";
        break;
      case CONNECTION_OK:
        s = "OK";
        break;
      case CONNECTION_AUTH_OK:
        s = "authOK";
        break;
      case CONNECTION_SSL_STARTUP:
        s = "sslStartup";
        break;
      case CONNECTION_SETENV:
        s = "setEnv";
        break;
      case CONNECTION_BAD:
        s = "bad";
        break;
      case CONNECTION_NEEDED:
        s = "needed";
        break;
    }

    return scope.Close(String::NewSymbol(s));
  }

  Connection () : EventEmitter () 
  {
    connection_ = NULL;

    connecting_ = resetting_ = false;

    ev_init(&read_watcher_, io_event);
    read_watcher_.data = this;

    ev_init(&write_watcher_, io_event);
    write_watcher_.data = this;
  }

  ~Connection ()
  {
    assert(connection_ == NULL);
  }

 private:

  void MakeConnection ()
  {
    PostgresPollingStatusType status;
    if (connecting_) {
     status = PQconnectPoll(connection_);
    } else {
     assert(resetting_);
     status = PQresetPoll(connection_);
    }

    if (status == PGRES_POLLING_READING) {
      ev_io_stop(EV_DEFAULT_ &write_watcher_);
      ev_io_start(EV_DEFAULT_ &read_watcher_);
      return;

    } else if (status == PGRES_POLLING_WRITING) {
      ev_io_stop(EV_DEFAULT_ &read_watcher_);
      ev_io_start(EV_DEFAULT_ &write_watcher_);
      return;
    }

    if (status == PGRES_POLLING_OK) {
      Emit(connect_symbol, 0, NULL);
      connecting_ = resetting_ = false;
      ev_io_start(EV_DEFAULT_ &read_watcher_);
      return;
    }

    CloseConnectionWithError();
  }

  Local<Value> BuildCell (PGresult *result, int row, int col)
  {
    HandleScope scope;

    if (PQgetisnull(result, row, col)) return scope.Close(Null());

    char *string = PQgetvalue(result, row, col);
    int32_t n = 0, i = 0;

    Oid t = PQftype(result, col);

    Handle<Value> cell;

    switch (t) {
      case BOOLOID:
        cell = *string == 't' ? True() : False();
        break;

      case INT8OID:
      case INT4OID:
      case INT2OID:
        for (i = string[0] == '-' ? 1 : 0; string[i]; i++) {
          n *= 10;
          n += string[i] - '0';
        }
        if (string[0] == '-') n = -n;
        cell = Integer::New(n);
        break;

      default:
#ifndef NDEBUG
//        printf("Unhandled OID: %d\n", t);
#endif
        cell = String::New(string);
    }
    return scope.Close(cell); 
  }

  Local<Value> BuildTuples (PGresult *result)
  {
    HandleScope scope;

    int nrows = PQntuples(result);
    int ncols = PQnfields(result);
    int row_index, col_index;
    char *field_name;

    Local<Array> tuples = Array::New(nrows);

    if (mapTupleItems_) {
      for (row_index = 0; row_index < nrows; row_index++) {
        Local<Object> row = Object::New();
        tuples->Set(Integer::New(row_index), row);

        for (col_index = 0; col_index < ncols; col_index++) {
          field_name = PQfname(result, col_index);
          Local<Value> cell = BuildCell(result, row_index, col_index); 
          row->Set(String::New(field_name), cell);
        }
      }
    } else {
      for (row_index = 0; row_index < nrows; row_index++) {
        Local<Array> row = Array::New(nrows);
        tuples->Set(Integer::New(row_index), row);

        for (col_index = 0; col_index < ncols; col_index++) {
          Local<Value> cell = BuildCell(result, row_index, col_index); 
          row->Set(Integer::New(col_index), cell);
        }
      }
    }

    return scope.Close(tuples);
  }

  void CloseConnectionWithError (const char *message_s = NULL)
  {
    HandleScope scope;

    if (!message_s) message_s = PQerrorMessage(connection_);
    Local<String> message = String::New(message_s);
    Local<Value> exception = Exception::Error(message);

    Close(exception);
  }

  static inline
  Local<Value> BuildResultException (PGresult *result)
  {
    HandleScope scope;

    char *primary_s = PQresultErrorField(result, PG_DIAG_MESSAGE_PRIMARY);
    assert(primary_s);
    Local<String> primary = String::New(primary_s);

    Local<Value> error_v = Exception::Error(primary);
    assert(error_v->IsObject());
    Local<Object> error = Local<Object>::Cast(error_v);

    char *full_s = PQresultErrorMessage(result);
    if (full_s) {
      error->Set(String::NewSymbol("full"), String::New(full_s));
    }

    char *detail_s = PQresultErrorField(result, PG_DIAG_MESSAGE_DETAIL);
    if (detail_s) {
      error->Set(String::NewSymbol("detail"), String::New(detail_s));
    }

    char *severity_s = PQresultErrorField(result, PG_DIAG_SEVERITY);
    if (severity_s) {
      error->Set(String::NewSymbol("severity"), String::New(severity_s));
    }

    /* TODO PG_DIAG_SQLSTATE
     *      PG_DIAG_MESSAGE_HINT
     *      PG_DIAG_STATEMENT_POSITION
     *      PG_DIAG_INTERNAL_POSITION
     *      PG_DIAG_INTERNAL_QUERY
     *      PG_DIAG_CONTEXT
     *      PG_DIAG_SOURCE_FILE
     *      PG_DIAG_SOURCE_LINE
     *      PG_DIAG_SOURCE_FUNCTION
     */

    return scope.Close(error);
  }

  void EmitResult (PGresult *result)
  {
    HandleScope scope;

    Local<Value> tuples;
    Local<Value> exception;

    Local<Value> args[2];

    switch (PQresultStatus(result)) {
      case PGRES_EMPTY_QUERY:
      case PGRES_COMMAND_OK:
        Emit(result_symbol, 0, NULL);
        break;

      case PGRES_TUPLES_OK:
        args[0] = Local<Value>::New(Null());
        args[1] = BuildTuples(result);
        Emit(result_symbol, 2, args);
        break;

      case PGRES_COPY_OUT:
      case PGRES_COPY_IN:
        assert(0 && "Not yet implemented.");
        exception = Exception::Error(String::New("Not yet implemented"));
        Emit(result_symbol, 1, &exception);
        break;

      case PGRES_BAD_RESPONSE:
      case PGRES_NONFATAL_ERROR:
      case PGRES_FATAL_ERROR:
        exception = BuildResultException(result);
        Emit(result_symbol, 1, &exception);
        break;
    }
  }

  void EmitNotify (PGnotify *notify)
  {
    HandleScope scope;

    Local<Value> args[3];

    args[0] = Local<Value>::New(String::New(notify->relname));
    args[1] = Local<Value>::New(Integer::New(notify->be_pid));
    args[2] = Local<Value>::New(String::New(notify->extra));
    Emit(notify_symbol, 3, args);
  }

  void Event (int revents)
  {
    if (revents & EV_ERROR) {
      CloseConnectionWithError("connection closed");
      return;
    }

    assert(PQisnonblocking(connection_));

    if (connecting_ || resetting_) {
      MakeConnection();
      return;
    }

    if (revents & EV_READ) {
      if (PQconsumeInput(connection_) == 0) {
        CloseConnectionWithError();
        return;
      }

      if (!PQisBusy(connection_)) {
        PGresult *result;
        while ((result = PQgetResult(connection_))) {
          EmitResult(result);
          PQclear(result);
        }
        Emit(ready_symbol, 0, NULL);
      }

      PGnotify *notify;
      while ((notify = PQnotifies(connection_))) {
        EmitNotify(notify);
        PQfreemem(notify);
      }
    }

    if (revents & EV_WRITE) {
      if (PQflush(connection_) == 0) ev_io_stop(EV_DEFAULT_ &write_watcher_);
    }
  }

  static void
  io_event (EV_P_ ev_io *w, int revents)
  {
    Connection *connection = static_cast<Connection*>(w->data);
    connection->Event(revents);
  }

  ev_io read_watcher_;
  ev_io write_watcher_;
  PGconn *connection_;
  bool connecting_;
  bool resetting_;
  bool mapTupleItems_;
};

extern "C" void
init (Handle<Object> target) 
{
  HandleScope scope;
  Connection::Initialize(target);
}