From d3bf5c20acda2fc1e635a0dad4bf5c88c110f567 Mon Sep 17 00:00:00 2001
From: Raymond Feng <raymond@strongloop.com>
Date: Thu, 8 Aug 2013 08:30:26 -0700
Subject: [PATCH] Add an optional models argument to automigrate/autoupdate

---
 lib/datasource.js | 28 ++++++++++++++++++++--------
 lib/sql.js        | 37 +++++++++++++++++++++++++++----------
 2 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/lib/datasource.js b/lib/datasource.js
index fde1e7ce..b83fa6e5 100644
--- a/lib/datasource.js
+++ b/lib/datasource.js
@@ -412,27 +412,39 @@ DataSource.prototype.defineProperty = function (model, prop, params) {
  * Drop each model table and re-create.
  * This method make sense only for sql connectors.
  *
+ * @param {String} or {[String]} Models to be migrated, if not present, apply to all models
+ *
  * @warning All data will be lost! Use autoupdate if you need your data.
  */
-DataSource.prototype.automigrate = function (cb) {
+DataSource.prototype.automigrate = function (models, cb) {
     this.freeze();
     if (this.connector.automigrate) {
-        this.connector.automigrate(cb);
-    } else if (cb) {
-        cb();
+        this.connector.automigrate(models, cb);
+    } else {
+        if ((!cb) && ('function' === typeof models)) {
+            cb = models;
+            models = undefined;
+        }
+        cb && process.nextTick(cb);
     }
 };
 
 /**
  * Update existing database tables.
  * This method make sense only for sql connectors.
+ *
+ * @param {String} or {[String]} Models to be migrated, if not present, apply to all models
  */
-DataSource.prototype.autoupdate = function (cb) {
+DataSource.prototype.autoupdate = function (models, cb) {
     this.freeze();
     if (this.connector.autoupdate) {
-        this.connector.autoupdate(cb);
-    } else if (cb) {
-        cb();
+        this.connector.autoupdate(models, cb);
+    } else {
+        if ((!cb) && ('function' === typeof models)) {
+            cb = models;
+            models = undefined;
+        }
+        cb && process.nextTick(cb);
     }
 };
 
diff --git a/lib/sql.js b/lib/sql.js
index 1817b19d..fdd932db 100644
--- a/lib/sql.js
+++ b/lib/sql.js
@@ -267,19 +267,36 @@ BaseSQL.prototype.disconnect = function disconnect() {
     this.client.end();
 };
 
-BaseSQL.prototype.automigrate = function (cb) {
+/**
+ * Recreate the tables for the given models
+ * @param models A model name or an array of model names, if not present, apply to all models defined in the connector
+ * @param cb Callback function
+ */
+BaseSQL.prototype.automigrate = function (models, cb) {
     var self = this;
     var wait = 0;
-    Object.keys(this._models).forEach(function (model) {
-        wait += 1;
-        self.dropTable(model, function () {
-            // console.log('drop', model);
-            self.createTable(model, function (err) {
-                // console.log('create', model);
-                if (err) console.log(err);
-                done();
+    if ((!cb) && ('function' === typeof models)) {
+        cb = models;
+        models = undefined;
+    }
+    // First argument is a model name
+    if ('string' === typeof models) {
+        models = [models];
+    }
+
+    models = models || Object.keys(this._models);
+    models.forEach(function (model) {
+        if (model in self._models) {
+            wait++;
+            self.dropTable(model, function () {
+                // console.log('drop', model);
+                self.createTable(model, function (err) {
+                    // console.log('create', model);
+                    if (err) console.log(err);
+                    done();
+                });
             });
-        });
+        }
     });
     if (wait === 0) cb();