From 3e2c728bdbee29ef018defcd65e898c3c06aec42 Mon Sep 17 00:00:00 2001
From: arily <arily@arily.moe>
Date: Fri, 23 Oct 2020 16:07:05 +0900
Subject: [PATCH] feat(mongo): add connection string support (#119)

---
 packages/plugin-mongo/src/database.ts | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/packages/plugin-mongo/src/database.ts b/packages/plugin-mongo/src/database.ts
index 7440ff4343..d0a42172f7 100644
--- a/packages/plugin-mongo/src/database.ts
+++ b/packages/plugin-mongo/src/database.ts
@@ -1,5 +1,6 @@
 import { MongoClient, Db, Collection } from 'mongodb'
 import { App, TableType } from 'koishi-core'
+import { URLSearchParams } from 'url'
 
 export interface Config {
   username?: string
@@ -7,8 +8,14 @@ export interface Config {
   protocol?: string
   host?: string
   port?: number
+  /** database name */
   name?: string
   prefix?: string
+  /** default auth database */
+  authDatabase?: string
+  connectOptions?: ConstructorParameters<typeof URLSearchParams>[0]
+  /** connection string (will overwrite all configs except 'name' and 'prefix') */
+  uri?: string
 }
 
 export default class MongoDatabase {
@@ -24,9 +31,7 @@ export default class MongoDatabase {
   }
 
   async start() {
-    let mongourl = `${this.config.protocol}://`
-    if (this.config.username) mongourl += `${this.config.username}:${this.config.password}@`
-    mongourl += `${this.config.host}:${this.config.port}/${this.config.name}`
+    const mongourl = this.config.uri || this.connectionStringFromConfig()
     this.client = await MongoClient.connect(
       mongourl, { useNewUrlParser: true, useUnifiedTopology: true },
     )
@@ -43,4 +48,16 @@ export default class MongoDatabase {
   stop() {
     return this.client.close()
   }
+
+  connectionStringFromConfig() {
+    const { authDatabase, connectOptions, host, name, password, port, protocol, username } = this.config
+    let mongourl = `${protocol}://`
+    if (username) mongourl += `${username}${password ? `:${password}` : ''}@`
+    mongourl += `${host}${port ? `:${port}` : ''}/${authDatabase || name}`
+    if (connectOptions) {
+      const params = new URLSearchParams(connectOptions)
+      mongourl += `?${params}`
+    }
+    return mongourl
+  }
 }