From 0a0fa3c8077059babf48656f5ed2c99d1a98f9f1 Mon Sep 17 00:00:00 2001
From: Dimas Restu H <dimas.restu@student.upi.edu>
Date: Mon, 25 Apr 2022 02:38:38 +0700
Subject: [PATCH] add capabilities to store in any database type

---
 internal/startup.go           |  2 +-
 internal/whatsapp/whatsapp.go |  6 +++-
 pkg/whatsapp/whatsapp.go      | 55 ++++++++++++++++++++++++-----------
 3 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/internal/startup.go b/internal/startup.go
index c5a1cb2..3feae45 100644
--- a/internal/startup.go
+++ b/internal/startup.go
@@ -23,7 +23,7 @@ func Startup() {
 
 		log.Print(nil).Info("Restoring WhatsApp Client for " + maskJID)
 
-		err := pkgWhatsApp.WhatAppConnect(jid)
+		err := pkgWhatsApp.WhatsAppInitClient(jid)
 		if err != nil {
 			log.Print(nil).Error(err.Error())
 		}
diff --git a/internal/whatsapp/whatsapp.go b/internal/whatsapp/whatsapp.go
index 75b6480..af23308 100644
--- a/internal/whatsapp/whatsapp.go
+++ b/internal/whatsapp/whatsapp.go
@@ -124,16 +124,20 @@ func Login(c echo.Context) error {
 		reqLogin.Output = "html"
 	}
 
-	err = pkgWhatsApp.WhatAppConnect(jid)
+	// Initialize WhatsApp Client
+	err = pkgWhatsApp.WhatsAppInitClient(jid)
 	if err != nil {
 		return router.ResponseInternalError(c, err.Error())
 	}
 
+	// Get WhatsApp QR Code Image
 	qrCodeImage, qrCodeTimeout, err := pkgWhatsApp.WhatsAppLogin(jid)
 	if err != nil {
 		return router.ResponseInternalError(c, err.Error())
 	}
 
+	// If Return is Not QR Code But Reconnected
+	// Then Return OK With Reconnected Status
 	if qrCodeImage == "WhatsApp Client is Reconnected" {
 		return router.ResponseSuccess(c, qrCodeImage)
 	}
diff --git a/pkg/whatsapp/whatsapp.go b/pkg/whatsapp/whatsapp.go
index 26d4732..26bfaa9 100644
--- a/pkg/whatsapp/whatsapp.go
+++ b/pkg/whatsapp/whatsapp.go
@@ -5,7 +5,6 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
-	"os"
 	"strings"
 
 	_ "github.com/mattn/go-sqlite3"
@@ -21,21 +20,39 @@ import (
 
 var WhatsAppClient = make(map[string]*whatsmeow.Client)
 
-func WhatsAppInit(jid string) (*whatsmeow.Client, error) {
-	// Prepare SQLite Database File and Connection Address
-	dbFileName := "dbs/" + jid + ".db"
-	dbAddress := fmt.Sprintf("file:%s?_foreign_keys=on", dbFileName)
+func WhatsAppDB(dbType string) (*sqlstore.Container, error) {
+	var dbName, dbURI string
 
-	// Create and Connect to SQLite Database
-	datastore, err := sqlstore.New("sqlite3", dbAddress, nil)
+	switch dbType {
+	case "sqlite3":
+		// Prepare SQLite Database and Connection URI
+		dbName = "dbs/WhatsApp.db"
+		dbURI = fmt.Sprintf("file:%s?_foreign_keys=on", dbName)
+
+	default:
+		return nil, errors.New("Unknown WhstaApp Client Database Type")
+	}
+
+	// Create and Connect to Database
+	datastore, err := sqlstore.New(dbType, dbURI, nil)
+	if err != nil {
+		return nil, errors.New("Failed to Connect WhatsApp Client Database")
+	}
+
+	return datastore, nil
+}
+
+func WhatsAppInitDB(jid string) (*whatsmeow.Client, error) {
+	// Connect to WhatsApp Client Datastore
+	datastore, err := WhatsAppDB("sqlite3")
 	if err != nil {
-		return nil, errors.New("Failed to Connect SQLite Database")
+		return nil, err
 	}
 
-	// Get First WhatsApp Device from SQLite Database
-	device, err := datastore.GetFirstDevice()
+	// Get WhatsApp Device Based on JID from Datastore
+	device, err := datastore.GetDevice(WhatsAppComposeJID(jid))
 	if err != nil {
-		return nil, errors.New("Failed to Load WhatsApp Device")
+		return nil, errors.New("Failed to Load WhatsApp Client Device from Database")
 	}
 
 	// Set Client Properties
@@ -49,10 +66,10 @@ func WhatsAppInit(jid string) (*whatsmeow.Client, error) {
 	return client, nil
 }
 
-func WhatAppConnect(jid string) error {
+func WhatsAppInitClient(jid string) error {
 	if WhatsAppClient[jid] == nil {
 		// Initialize New WhatsApp Client
-		client, err := WhatsAppInit(jid)
+		client, err := WhatsAppInitDB(jid)
 		if err != nil {
 			return err
 		}
@@ -156,11 +173,15 @@ func WhatsAppLogout(jid string) error {
 		// Logout WhatsApp Client and Disconnect from WebSocket
 		err := WhatsAppClient[jid].Logout()
 		if err != nil {
-			return err
-		}
+			// Force Disconnect
+			WhatsAppClient[jid].Disconnect()
 
-		// Remove SQLite Database File
-		_ = os.Remove("dbs/" + jid + ".db")
+			// Manually Delete Device from Database Store
+			err = WhatsAppClient[jid].Store.Delete()
+			if err != nil {
+				return err
+			}
+		}
 
 		// Free WhatsApp Client Map
 		WhatsAppClient[jid] = nil