From b8c56230410b0f60c836edb16ec534831df813ae Mon Sep 17 00:00:00 2001
From: Noah Zalev <github@zalev.ca>
Date: Sat, 1 May 2021 18:41:28 -0400
Subject: [PATCH] Bugfix - don't rely on use database

---
 db.js | 80 +++++++++++++++++++++++++++++++++++++----------------------
 1 file changed, 50 insertions(+), 30 deletions(-)

diff --git a/db.js b/db.js
index f90d70f6..00b67a82 100644
--- a/db.js
+++ b/db.js
@@ -438,37 +438,8 @@ module.exports.CreateDB = function (parent, func) {
         obj.dbRecordsDecryptKey = parent.crypto.createHash('sha384').update(parent.args.dbrecordsdecryptkey).digest('raw').slice(0, 32);
     }
 
-    if (parent.args.mariadb || parent.args.mysql) {
-        var connectinArgs = (parent.args.mariadb) ? parent.args.mariadb : parent.args.mysql;
-        var dbname = (connectinArgs.database != null) ? connectinArgs.database : 'meshcentral';
 
-        // Including the db name in the connection obj will cause a connection failure if it does not exist
-        var connectionObject = Clone(connectinArgs);
-        delete connectionObject.database;
-
-        try {
-            if (connectinArgs.ssl) {
-                if (connectinArgs.ssl.cacertpath) { connectionObject.ssl.ca = [require('fs').readFileSync(connectinArgs.ssl.cacertpath, 'utf8')]; }
-                if (connectinArgs.ssl.clientcertpath) { connectionObject.ssl.cert = [require('fs').readFileSync(connectinArgs.ssl.clientcertpath, 'utf8')]; }
-                if (connectinArgs.ssl.clientkeypath) { connectionObject.ssl.key = [require('fs').readFileSync(connectinArgs.ssl.clientkeypath, 'utf8')]; }
-            }
-        } catch (ex) {
-            console.log('Error loading SQL Connector certificate: ' + ex);
-            process.exit();
-        }
-
-        if (parent.args.mariadb) {
-            // Use MariaDB
-            obj.databaseType = 4;
-            Datastore = require('mariadb').createPool(connectionObject);
-        } else if (parent.args.mysql) {
-            // Use MySQL
-            Datastore = require('mysql').createConnection(connectionObject);
-            obj.databaseType = 5;
-        }
-        sqlDbQuery('CREATE DATABASE IF NOT EXISTS ' + dbname);
-
-        // Set the default database for the rest of this connections lifetime
+    function createTablesIfNotExist(dbname) {
         var useDatabase = 'USE ' + dbname;
         sqlDbQuery(useDatabase, null, function (err, docs) {
             if (err != null) {
@@ -503,6 +474,55 @@ module.exports.CreateDB = function (parent, func) {
                 });
             }
         });
+    }
+
+    if (parent.args.mariadb || parent.args.mysql) {
+        var connectinArgs = (parent.args.mariadb) ? parent.args.mariadb : parent.args.mysql;
+        var dbname = (connectinArgs.database != null) ? connectinArgs.database : 'meshcentral';
+
+        // Including the db name in the connection obj will cause a connection failure if it does not exist
+        var connectionObject = Clone(connectinArgs);
+        delete connectionObject.database;
+
+        try {
+            if (connectinArgs.ssl) {
+                if (connectinArgs.ssl.cacertpath) { connectionObject.ssl.ca = [require('fs').readFileSync(connectinArgs.ssl.cacertpath, 'utf8')]; }
+                if (connectinArgs.ssl.clientcertpath) { connectionObject.ssl.cert = [require('fs').readFileSync(connectinArgs.ssl.clientcertpath, 'utf8')]; }
+                if (connectinArgs.ssl.clientkeypath) { connectionObject.ssl.key = [require('fs').readFileSync(connectinArgs.ssl.clientkeypath, 'utf8')]; }
+            }
+        } catch (ex) {
+            console.log('Error loading SQL Connector certificate: ' + ex);
+            process.exit();
+        }
+
+        if (parent.args.mariadb) {
+            // Use MariaDB
+            obj.databaseType = 4;
+            var tempDatastore = require('mariadb').createPool(connectionObject);
+            tempDatastore.getConnection().then(function (conn) {
+                conn.query('CREATE DATABASE IF NOT EXISTS ' + dbname).then(function (result) {
+                    conn.release();
+                }).catch(function (ex) { console.log('Auto-create database failed: ' + ex); });
+            }).catch(function (ex) { console.log('Auto-create database failed: ' + ex); });
+            setTimeout(function () { tempDatastore.end(); }, 2000);
+
+            connectionObject.database = dbname;
+            Datastore = require('mariadb').createPool(connectionObject);
+            createTablesIfNotExist(dbname);
+        } else if (parent.args.mysql) {
+            // Use MySQL
+            obj.databaseType = 5;
+            var tempDatastore = require('mysql').createConnection(connectionObject);
+            tempDatastore.query('CREATE DATABASE IF NOT EXISTS ' + dbname, function (error) {
+                if (error != null) {
+                    console.log('Auto-create database failed: ' + error);
+                }
+                connectionObject.database = dbname;
+                Datastore = require('mysql').createConnection(connectionObject);
+                createTablesIfNotExist(dbname);
+            });
+            setTimeout(function () { tempDatastore.end(); }, 2000);
+        }
     } else if (parent.args.mongodb) {
         // Use MongoDB
         obj.databaseType = 3;