From 5fa5f212897cd523d2994a8c2559493ea02b875a Mon Sep 17 00:00:00 2001 From: TheEssem Date: Fri, 18 Dec 2020 18:50:25 -0600 Subject: [PATCH] Moved database driver-specific code to their own files --- app.js | 2 +- events/ready.js | 2 +- utils/convertdb.js | 5 +- utils/database.js | 247 +------------------------------------ utils/database/mongo.js | 137 ++++++++++++++++++++ utils/database/postgres.js | 107 ++++++++++++++++ 6 files changed, 249 insertions(+), 251 deletions(-) create mode 100644 utils/database/mongo.js create mode 100644 utils/database/postgres.js diff --git a/app.js b/app.js index 4ba3e6f..69045b3 100644 --- a/app.js +++ b/app.js @@ -61,7 +61,7 @@ async function init() { client.disconnect(); const db = require("./utils/database.js"); if (process.env.DB_DRIVER=== "mongo") { - db.connection.close(() => { + db.connection.disconnect(() => { process.exit(0); }); } else if (process.env.DB_DRIVER=== "postgres") { diff --git a/events/ready.js b/events/ready.js index 6c24998..5ff140d 100644 --- a/events/ready.js +++ b/events/ready.js @@ -15,7 +15,7 @@ module.exports = async () => { // connect to lavalink if (!soundPlayer.status && !soundPlayer.connected) await soundPlayer.connect(); - await database.handleCounts(); + await database.setup(); // generate docs if (helpGenerator && first) await helpGenerator(process.env.OUTPUT); diff --git a/utils/convertdb.js b/utils/convertdb.js index f4cbb85..0308f8e 100644 --- a/utils/convertdb.js +++ b/utils/convertdb.js @@ -1,10 +1,7 @@ require("dotenv").config(); const { Pool } = require("pg"); const pool = new Pool({ - user: "esmbot", - host: "localhost", - database: "esmbot", - port: 5432 + connectionString: process.env.DB }); const mongoose = require("mongoose"); mongoose.connect(process.env.MONGO, { poolSize: 10, bufferMaxEntries: 0, useNewUrlParser: true, useUnifiedTopology: true }); diff --git a/utils/database.js b/utils/database.js index 6713c40..51c1d2a 100644 --- a/utils/database.js +++ b/utils/database.js @@ -1,246 +1,3 @@ -// database stuff -const logger = require("./logger.js"); -const collections = require("../utils/collections.js"); -const misc = require("./misc.js"); +// wrapper for the database drivers in ./database/ -if (process.env.DB_DRIVER=== "mongo") { - const mongoose = require("mongoose"); - mongoose.connect(process.env.DB, { - poolSize: 10, - bufferMaxEntries: 0, - useNewUrlParser: true, - useUnifiedTopology: true, - }); - const guildSchema = new mongoose.Schema({ - id: String, - tags: Map, - prefix: String, - disabled: [String], - tagsDisabled: Boolean - }); - const Guild = mongoose.model("Guild", guildSchema); - - const globalSchema = new mongoose.Schema({ - cmdCounts: Map, - }); - const Global = mongoose.model("Global", globalSchema); - - exports.guilds = Guild; - exports.global = Global; - exports.connection = mongoose.connection; -} else if (process.env.DB_DRIVER=== "postgres") { - const { Pool } = require("pg"); - const pool = new Pool({ - connectionString: process.env.DB - }); - exports.connection = pool; -} - -exports.getGuild = async (query) => { - if (process.env.DB_DRIVER=== "mongo") { - return await this.guilds.findOne({ id: query }); - } else if (process.env.DB_DRIVER=== "postgres") { - return (await this.connection.query("SELECT * FROM guilds WHERE guild_id = $1", [query])).rows[0]; - } -}; - -exports.setPrefix = async (prefix, guild) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = await this.getGuild(guild.id); - guildDB.prefix = prefix; - await guildDB.save(); - collections.prefixCache.set(guild.id, prefix); - } else if (process.env.DB_DRIVER=== "postgres") { - await this.connection.query("UPDATE guilds SET prefix = $1 WHERE guild_id = $2", [prefix, guild.id]); - collections.prefixCache.set(guild.id, prefix); - } -}; - -exports.setTag = async (name, content, guild) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = await this.getGuild(guild.id); - guildDB.tags[name] = content; - await guildDB.save(); - } else if (process.env.DB_DRIVER=== "postgres") { - const guildDB = await this.getGuild(guild.id); - guildDB.tags[name] = content; - await this.connection.query("UPDATE guilds SET tags = $1 WHERE guild_id = $2", [guildDB.tags, guild.id]); - } -}; - -exports.removeTag = async (name, guild) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = await this.getGuild(guild.id); - delete guildDB.tags[name]; - await guildDB.save(); - } else if (process.env.DB_DRIVER=== "postgres") { - const guildDB = await this.getGuild(guild.id); - delete guildDB.tags[name]; - await this.connection.query("UPDATE guilds SET tags = $1 WHERE guild_id = $2", [guildDB.tags, guild.id]); - } -}; - -exports.toggleTags = async (guild) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = await this.getGuild(guild.id); - guildDB.tagsDisabled = !guildDB.tagsDisabled; - await guildDB.save(); - return guildDB.tagsDisabled; - } else if (process.env.DB_DRIVER=== "postgres") { - const guildDB = await this.getGuild(guild.id); - guildDB.tags_disabled = !guildDB.tags_disabled; - await this.connection.query("UPDATE guilds SET tags_disabled = $1 WHERE guild_id = $2", [guildDB.tags_disabled, guild.id]); - return guildDB.tags_disabled; - } -}; - -exports.disableChannel = async (channel) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = await this.getGuild(channel.guild.id); - guildDB.disabled.push(channel.id); - await guildDB.save(); - collections.disabledCache.set(channel.guild.id, guildDB.disabled); - } else if (process.env.DB_DRIVER=== "postgres") { - const guildDB = await this.getGuild(channel.guild.id); - await this.connection.query("UPDATE guilds SET disabled = $1 WHERE guild_id = $2", [[...guildDB.disabled, channel.id], channel.guild.id]); - collections.disabledCache.set(channel.guild.id, guildDB.disabled); - } -}; - -exports.enableChannel = async (channel) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = await this.getGuild(channel.guild.id); - guildDB.disabled = guildDB.disabled.filter(item => item !== channel.id); - await guildDB.save(); - collections.disabledCache.set(channel.guild.id, guildDB.disabled); - } else if (process.env.DB_DRIVER=== "postgres") { - const guildDB = await this.getGuild(channel.guild.id); - const newDisabled = guildDB.disabled.filter(item => item !== channel.id); - await this.connection.query("UPDATE guilds SET disabled = $1 WHERE guild_id = $2", [newDisabled, channel.guild.id]); - collections.disabledCache.set(channel.guild.id, guildDB.disabled); - } -}; - -exports.getCounts = async () => { - if (process.env.DB_DRIVER=== "mongo") { - return [...(await this.global.findOne({})).cmdCounts.entries()]; - } else if (process.env.DB_DRIVER=== "postgres") { - const counts = await this.connection.query("SELECT * FROM counts"); - const countArray = []; - for (const { command, count } of counts.rows) { - countArray.push([command, count]); - } - return countArray; - } -}; - -exports.addCount = async (command) => { - if (process.env.DB_DRIVER=== "mongo") { - const global = await this.global.findOne({}); - const count = global.cmdCounts.get(command); - global.cmdCounts.set(command, parseInt(count) + 1); - await global.save(); - } else if (process.env.DB_DRIVER=== "postgres") { - const count = await this.connection.query("SELECT * FROM counts WHERE command = $1", [command]); - await this.connection.query("UPDATE counts SET count = $1 WHERE command = $2", [count.rows[0].count + 1, command]); - } -}; - -exports.addGuild = async (guild) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = new this.guilds({ - id: guild.id, - tags: misc.tagDefaults, - prefix: process.env.PREFIX, - disabled: [], - tagsDisabled: false - }); - await guildDB.save(); - return guildDB; - } else if (process.env.DB_DRIVER=== "postgres") { - await this.connection.query("INSERT INTO guilds (guild_id, tags, prefix, warns, disabled, tags_disabled) VALUES ($1, $2, $3, $4, $5, $6)", [guild.id, misc.tagDefaults, process.env.PREFIX, {}, [], false]); - return await this.getGuild(guild.id); - } -}; - -exports.fixGuild = async (guild) => { - if (process.env.DB_DRIVER=== "mongo") { - const guildDB = await this.guilds.findOne({ id: guild.id }); - if (!guildDB) { - logger.log(`Registering guild database entry for guild ${guild.id}...`); - return await this.addGuild(guild); - } else { - if (!guildDB.disabled && guildDB.disabledChannels) { - guildDB.set("disabled", guildDB.disabledChannels); - guildDB.set("disabledChannels", undefined); - await guildDB.save(); - return guildDB; - } - } - } else if (process.env.DB_DRIVER=== "postgres") { - const guildDB = await this.connection.query("SELECT * FROM guilds WHERE guild_id = $1", [guild.id]); - if (guildDB.rows.length === 0) { - logger.log(`Registering guild database entry for guild ${guild.id}...`); - return await this.addGuild(guild); - } - } -}; - -exports.handleCounts = async () => { - if (process.env.DB_DRIVER=== "mongo") { - const global = await this.global.findOne({}); - if (!global) { - const countObject = {}; - for (const command of collections.commands.keys()) { - countObject[command] = 0; - } - const newGlobal = new this.global({ - cmdCounts: countObject - }); - await newGlobal.save(); - } else { - const exists = []; - for (const command of collections.commands.keys()) { - if (!global.cmdCounts.has(command)) { - global.cmdCounts.set(command, 0); - } - exists.push(command); - } - - for (const command of global.cmdCounts.keys()) { - if (!exists.includes(command)) { - global.cmdCounts.set(command, undefined); - } - } - await global.save(); - } - } else if (process.env.DB_DRIVER=== "postgres") { - let counts; - try { - counts = await this.connection.query("SELECT * FROM counts"); - } catch { - counts = { rows: [] }; - } - - if (!counts.rows[0]) { - for (const command of collections.commands.keys()) { - await this.connection.query("INSERT INTO counts (command, count) VALUES ($1, $2)", [command, 0]); - } - } else { - const exists = []; - for (const command of collections.commands.keys()) { - const count = await this.connection.query("SELECT * FROM counts WHERE command = $1", [command]); - if (!count.rows[0]) { - await this.connection.query("INSERT INTO counts (command, count) VALUES ($1, $2)", [command, 0]); - } - exists.push(command); - } - - for (const { command } of counts.rows) { - if (!exists.includes(command)) { - await this.connection.query("DELETE FROM counts WHERE command = $1", [command]); - } - } - } - } -}; \ No newline at end of file +module.exports = require(`./database/${process.env.DB_DRIVER}.js`); \ No newline at end of file diff --git a/utils/database/mongo.js b/utils/database/mongo.js new file mode 100644 index 0000000..e9ebaa3 --- /dev/null +++ b/utils/database/mongo.js @@ -0,0 +1,137 @@ +const collections = require("../collections.js"); +const logger = require("../logger.js"); +const misc = require("../misc.js"); + +const mongoose = require("mongoose"); +mongoose.connect(process.env.DB, { + poolSize: 10, + bufferMaxEntries: 0, + useNewUrlParser: true, + useUnifiedTopology: true, +}); +const guildSchema = new mongoose.Schema({ + id: String, + tags: Map, + prefix: String, + disabled: [String], + tagsDisabled: Boolean +}); +const Guild = mongoose.model("Guild", guildSchema); + +const globalSchema = new mongoose.Schema({ + cmdCounts: Map, +}); +const Global = mongoose.model("Global", globalSchema); + +exports.connection = mongoose.connection; + +exports.getGuild = async (query) => { + return await Guild.findOne({ id: query }); +}; + +exports.setPrefix = async (prefix, guild) => { + const guildDB = await this.getGuild(guild.id); + guildDB.prefix = prefix; + await guildDB.save(); + collections.prefixCache.set(guild.id, prefix); +}; + +exports.setTag = async (name, content, guild) => { + const guildDB = await this.getGuild(guild.id); + guildDB.tags[name] = content; + await guildDB.save(); +}; + +exports.removeTag = async (name, guild) => { + const guildDB = await this.getGuild(guild.id); + delete guildDB.tags[name]; + await guildDB.save(); +}; + +exports.toggleTags = async (guild) => { + const guildDB = await this.getGuild(guild.id); + guildDB.tagsDisabled = !guildDB.tagsDisabled; + await guildDB.save(); + return guildDB.tagsDisabled; +}; + +exports.disableChannel = async (channel) => { + const guildDB = await this.getGuild(channel.guild.id); + guildDB.disabled.push(channel.id); + await guildDB.save(); + collections.disabledCache.set(channel.guild.id, guildDB.disabled); +}; + +exports.enableChannel = async (channel) => { + const guildDB = await this.getGuild(channel.guild.id); + guildDB.disabled = guildDB.disabled.filter(item => item !== channel.id); + await guildDB.save(); + collections.disabledCache.set(channel.guild.id, guildDB.disabled); +}; + +exports.getCounts = async () => { + return [...(await Global.findOne({})).cmdCounts.entries()]; +}; + +exports.addCount = async (command) => { + const global = await Global.findOne({}); + const count = global.cmdCounts.get(command); + global.cmdCounts.set(command, parseInt(count) + 1); + await global.save(); +}; + +exports.addGuild = async (guild) => { + const guildDB = new Guild({ + id: guild.id, + tags: misc.tagDefaults, + prefix: process.env.PREFIX, + disabled: [], + tagsDisabled: false + }); + await guildDB.save(); + return guildDB; +}; + +exports.fixGuild = async (guild) => { + const guildDB = await Guild.findOne({ id: guild.id }); + if (!guildDB) { + logger.log(`Registering guild database entry for guild ${guild.id}...`); + return await this.addGuild(guild); + } else { + if (!guildDB.disabled && guildDB.disabledChannels) { + guildDB.set("disabled", guildDB.disabledChannels); + guildDB.set("disabledChannels", undefined); + await guildDB.save(); + return guildDB; + } + } +}; + +exports.setup = async () => { + const global = await Global.findOne({}); + if (!global) { + const countObject = {}; + for (const command of collections.commands.keys()) { + countObject[command] = 0; + } + const newGlobal = new Global({ + cmdCounts: countObject + }); + await newGlobal.save(); + } else { + const exists = []; + for (const command of collections.commands.keys()) { + if (!global.cmdCounts.has(command)) { + global.cmdCounts.set(command, 0); + } + exists.push(command); + } + + for (const command of global.cmdCounts.keys()) { + if (!exists.includes(command)) { + global.cmdCounts.set(command, undefined); + } + } + await global.save(); + } +}; \ No newline at end of file diff --git a/utils/database/postgres.js b/utils/database/postgres.js new file mode 100644 index 0000000..333e502 --- /dev/null +++ b/utils/database/postgres.js @@ -0,0 +1,107 @@ +const collections = require("../collections.js"); +const logger = require("../logger.js"); +const misc = require("../misc.js"); + +const { Pool } = require("pg"); +const pool = new Pool({ + connectionString: process.env.DB +}); +const connection = pool; + +exports.getGuild = async (query) => { + return (await connection.query("SELECT * FROM guilds WHERE guild_id = $1", [query])).rows[0]; +}; + +exports.setPrefix = async (prefix, guild) => { + await connection.query("UPDATE guilds SET prefix = $1 WHERE guild_id = $2", [prefix, guild.id]); + collections.prefixCache.set(guild.id, prefix); +}; + +exports.setTag = async (name, content, guild) => { + const guildDB = await this.getGuild(guild.id); + guildDB.tags[name] = content; + await connection.query("UPDATE guilds SET tags = $1 WHERE guild_id = $2", [guildDB.tags, guild.id]); +}; + +exports.removeTag = async (name, guild) => { + const guildDB = await this.getGuild(guild.id); + delete guildDB.tags[name]; + await connection.query("UPDATE guilds SET tags = $1 WHERE guild_id = $2", [guildDB.tags, guild.id]); +}; + +exports.toggleTags = async (guild) => { + const guildDB = await this.getGuild(guild.id); + guildDB.tags_disabled = !guildDB.tags_disabled; + await connection.query("UPDATE guilds SET tags_disabled = $1 WHERE guild_id = $2", [guildDB.tags_disabled, guild.id]); + return guildDB.tags_disabled; +}; + +exports.disableChannel = async (channel) => { + const guildDB = await this.getGuild(channel.guild.id); + await connection.query("UPDATE guilds SET disabled = $1 WHERE guild_id = $2", [[...guildDB.disabled, channel.id], channel.guild.id]); + collections.disabledCache.set(channel.guild.id, guildDB.disabled); +}; + +exports.enableChannel = async (channel) => { + const guildDB = await this.getGuild(channel.guild.id); + const newDisabled = guildDB.disabled.filter(item => item !== channel.id); + await connection.query("UPDATE guilds SET disabled = $1 WHERE guild_id = $2", [newDisabled, channel.guild.id]); + collections.disabledCache.set(channel.guild.id, guildDB.disabled); +}; + +exports.getCounts = async () => { + const counts = await connection.query("SELECT * FROM counts"); + const countArray = []; + for (const { command, count } of counts.rows) { + countArray.push([command, count]); + } + return countArray; +}; + +exports.addCount = async (command) => { + const count = await connection.query("SELECT * FROM counts WHERE command = $1", [command]); + await connection.query("UPDATE counts SET count = $1 WHERE command = $2", [count.rows[0].count + 1, command]); +}; + +exports.addGuild = async (guild) => { + await connection.query("INSERT INTO guilds (guild_id, tags, prefix, warns, disabled, tags_disabled) VALUES ($1, $2, $3, $4, $5, $6)", [guild.id, misc.tagDefaults, process.env.PREFIX, {}, [], false]); + return await this.getGuild(guild.id); +}; + +exports.fixGuild = async (guild) => { + const guildDB = await connection.query("SELECT * FROM guilds WHERE guild_id = $1", [guild.id]); + if (guildDB.rows.length === 0) { + logger.log(`Registering guild database entry for guild ${guild.id}...`); + return await this.addGuild(guild); + } +}; + +exports.setup = async () => { + let counts; + try { + counts = await connection.query("SELECT * FROM counts"); + } catch { + counts = { rows: [] }; + } + + if (!counts.rows[0]) { + for (const command of collections.commands.keys()) { + await connection.query("INSERT INTO counts (command, count) VALUES ($1, $2)", [command, 0]); + } + } else { + const exists = []; + for (const command of collections.commands.keys()) { + const count = await connection.query("SELECT * FROM counts WHERE command = $1", [command]); + if (!count.rows[0]) { + await connection.query("INSERT INTO counts (command, count) VALUES ($1, $2)", [command, 0]); + } + exists.push(command); + } + + for (const { command } of counts.rows) { + if (!exists.includes(command)) { + await connection.query("DELETE FROM counts WHERE command = $1", [command]); + } + } + } +}; \ No newline at end of file