From 788bd8e09cd3404d304302060a36d34bb08f3d39 Mon Sep 17 00:00:00 2001
From: Charles Pick <charles.pick@gmail.com>
Date: Mon, 24 Nov 2014 15:24:31 +0000
Subject: [PATCH] add db.registerTransformer()

---
 lib/db/index.js    | 30 +++++++++++++++++
 test/db/db-test.js | 81 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 111 insertions(+)

diff --git a/lib/db/index.js b/lib/db/index.js
index de0370d..6e0cdb3 100644
--- a/lib/db/index.js
+++ b/lib/db/index.js
@@ -52,6 +52,7 @@ Db.prototype.configure = function (config) {
   this.password = config.password || 'admin';
   this.dataSegments = [];
   this.transactionId = 0;
+  this.transformers = config.transformers || {};
   return this;
 };
 
@@ -302,6 +303,9 @@ Db.prototype.normalizeResultContent = function (content) {
       cluster: content.cluster,
       position: content.position
     });
+    if (value['@class']) {
+      return this.transformDocument(value);
+    }
     return value;
   }
   else {
@@ -309,6 +313,32 @@ Db.prototype.normalizeResultContent = function (content) {
   }
 };
 
+Db.prototype.registerTransformer = function (className, transformer) {
+  this.transformers[className] = this.transformers[className] || [];
+  this.transformers[className].push(transformer);
+  return this;
+};
+
+
+/**
+ * Transform a document according to its `@class` property, using the registered transformers.
+ * @param  {Object} document The document to transform.
+ * @return {Mixed}           The transformed document.
+ */
+Db.prototype.transformDocument = function (document) {
+  var className = document['@class'];
+  if (this.transformers[className]) {
+    return this.transformers[className].reduce(function (document, transformer) {
+      return transformer(document);
+    }, document);
+  }
+  else {
+    return document;
+  }
+};
+
+
+
 // # Query Builder Methods
 
 /**
diff --git a/test/db/db-test.js b/test/db/db-test.js
index f5c3855..6df7dbb 100644
--- a/test/db/db-test.js
+++ b/test/db/db-test.js
@@ -1,3 +1,5 @@
+var Promise = require('bluebird');
+
 describe("Database API", function () {
   before(function () {
     return TEST_SERVER.create({
@@ -72,4 +74,83 @@ describe("Database API", function () {
       });
     });
   });
+
+  describe('Db::registerTransformer()', function () {
+    function OUser (data) {
+      if (!(this instanceof OUser)) {
+        return new OUser(data);
+      }
+      var keys = Object.keys(data),
+          length = keys.length,
+          key, i;
+      for (i = 0; i < length; i++) {
+        key = keys[i];
+        this[key] = data[key];
+      }
+    }
+
+    function ORole (data) {
+      if (!(this instanceof ORole)) {
+        return new ORole(data);
+      }
+      var keys = Object.keys(data),
+          length = keys.length,
+          key, i;
+      for (i = 0; i < length; i++) {
+        key = keys[i];
+        this[key] = data[key];
+      }
+    }
+
+
+    before(function () {
+      this.db.registerTransformer('OUser', OUser);
+      this.db.registerTransformer('ORole', ORole);
+    });
+
+    it('should register a transformation function for a class', function () {
+      return Promise.all([
+        this.db.select().from('OUser').all(),
+        this.db.select().from('ORole').all()
+      ])
+      .spread(function (users, roles) {
+        users.length.should.be.above(0);
+        users.forEach(function (user) {
+          user.should.be.an.instanceOf(OUser);
+        });
+        roles.length.should.be.above(0);
+        roles.forEach(function (role) {
+          role.should.be.an.instanceOf(ORole);
+        });
+      });
+    });
+
+    it('should transform documents even when they are nested', function () {
+      return this.db.select().from('OUser').fetch({roles: 1}).all()
+      .then(function (users) {
+        users.length.should.be.above(0);
+        users.forEach(function (user) {
+          user.should.be.an.instanceOf(OUser);
+          user.roles.length.should.be.above(0);
+          user.roles.forEach(function (role) {
+            role.should.be.an.instanceOf(ORole);
+          });
+        });
+      });
+    });
+
+    it('should still allow scalars', function () {
+      return this.db.select().from('OUser').limit(1).scalar()
+      .then(function (result) {
+        result.should.equal('admin');
+      });
+    });
+
+    it('should not transform when individual columns are selected', function () {
+      return this.db.select('name, status').from('OUser').limit(1).one()
+      .then(function (result) {
+        result.should.not.be.an.instanceOf(OUser);
+      });
+    });
+  });
 });
\ No newline at end of file