From 191f00cc89005b8def923da4029a35e11684aa49 Mon Sep 17 00:00:00 2001
From: NGPixel <ngpixel@gmail.com>
Date: Thu, 30 Mar 2017 21:53:21 -0400
Subject: [PATCH] Allow ENV variable as MongoDB conn string

---
 agent.js          |  2 +-
 config.sample.yml |  1 +
 libs/db.js        | 70 +++++++++++++++++++++++++++++++++++++++++++++++
 server.js         |  2 +-
 4 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 libs/db.js

diff --git a/agent.js b/agent.js
index f4add906..f2461b7c 100644
--- a/agent.js
+++ b/agent.js
@@ -28,7 +28,7 @@ winston.info('[AGENT] Background Agent is initializing...')
 let appconf = require(CORE_PATH + 'core-libs/config')()
 global.appconfig = appconf.config
 global.appdata = appconf.data
-global.db = require(CORE_PATH + 'core-libs/mongodb').init()
+global.db = require('./libs/db').init()
 global.upl = require('./libs/uploads-agent').init()
 global.git = require('./libs/git').init()
 global.entries = require('./libs/entries').init()
diff --git a/config.sample.yml b/config.sample.yml
index a5b543bc..d5b923a7 100644
--- a/config.sample.yml
+++ b/config.sample.yml
@@ -105,6 +105,7 @@ admin: admin@company.com
 # ---------------------------------------------------------------------
 # Database Connection String
 # ---------------------------------------------------------------------
+# You can also use an ENV variable by using $ENV_VAR_NAME as the value
 
 db: mongodb://localhost:27017/wiki
 
diff --git a/libs/db.js b/libs/db.js
new file mode 100644
index 00000000..4ad4a40f
--- /dev/null
+++ b/libs/db.js
@@ -0,0 +1,70 @@
+'use strict'
+
+/* global ROOTPATH, appconfig, winston */
+
+const modb = require('mongoose')
+const fs = require('fs')
+const path = require('path')
+const _ = require('lodash')
+
+/**
+ * MongoDB module
+ *
+ * @return     {Object}  MongoDB wrapper instance
+ */
+module.exports = {
+
+  /**
+   * Initialize DB
+   *
+   * @return     {Object}  DB instance
+   */
+  init () {
+    let self = this
+    global.Mongoose = modb
+
+    let dbModelsPath = path.resolve(ROOTPATH, 'models')
+
+    modb.Promise = require('bluebird')
+
+    // Event handlers
+
+    modb.connection.on('error', err => {
+      winston.error('Failed to connect to MongoDB instance.')
+      return err
+    })
+    modb.connection.once('open', function () {
+      winston.log('Connected to MongoDB instance.')
+    })
+
+    // Store connection handle
+
+    self.connection = modb.connection
+    self.ObjectId = modb.Types.ObjectId
+
+    // Load DB Models
+
+    fs
+    .readdirSync(dbModelsPath)
+    .filter(function (file) {
+      return (file.indexOf('.') !== 0)
+    })
+    .forEach(function (file) {
+      let modelName = _.upperFirst(_.camelCase(_.split(file, '.')[0]))
+      self[modelName] = require(path.join(dbModelsPath, file))
+    })
+
+    // Using ENV variable?
+
+    if (_.startsWith(appconfig.db, '$')) {
+      appconfig.db = process.env[appconfig.db.slice(1)]
+    }
+
+    // Connect
+
+    self.onReady = modb.connect(appconfig.db)
+
+    return self
+  }
+
+}
diff --git a/server.js b/server.js
index 977504c0..0076551f 100644
--- a/server.js
+++ b/server.js
@@ -33,7 +33,7 @@ winston.info('[SERVER] Wiki.js is initializing...')
 // ----------------------------------------
 
 global.lcdata = require('./libs/local').init()
-global.db = require(CORE_PATH + 'core-libs/mongodb').init()
+global.db = require('./libs/db').init()
 global.entries = require('./libs/entries').init()
 global.git = require('./libs/git').init(false)
 global.lang = require('i18next')