From b5802adf5beb6c638ce993f07f9ba6adc0fe66e2 Mon Sep 17 00:00:00 2001
From: harsh-98 <harshjniitr@gmail.com>
Date: Thu, 5 Oct 2023 07:20:02 +0700
Subject: [PATCH] fix: use NewQueries from db utils

---
 cmd/waku/node.go                |  3 +--
 waku/persistence/driver_type.go | 22 ++++++++++++++++++++++
 waku/persistence/store.go       |  8 ++++----
 waku/persistence/utils/db.go    | 11 +++++++++++
 4 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 waku/persistence/driver_type.go

diff --git a/cmd/waku/node.go b/cmd/waku/node.go
index 0f5e6873..7edef337 100644
--- a/cmd/waku/node.go
+++ b/cmd/waku/node.go
@@ -18,7 +18,6 @@ import (
 	"github.com/pbnjay/memory"
 	"github.com/prometheus/client_golang/prometheus"
 
-	"github.com/waku-org/go-waku/waku/persistence/sqlite"
 	dbutils "github.com/waku-org/go-waku/waku/persistence/utils"
 	wakupeerstore "github.com/waku-org/go-waku/waku/v2/peerstore"
 	"github.com/waku-org/go-waku/waku/v2/rendezvous"
@@ -203,7 +202,7 @@ func Execute(options NodeOptions) {
 
 	if options.Store.Enable && options.PersistPeers {
 		// Create persistent peerstore
-		queries, err := sqlite.NewQueries("peerstore", db)
+		queries, err := dbutils.NewQueries("peerstore", db)
 		failOnErr(err, "Peerstore")
 
 		datastore := dssql.NewDatastore(db, queries)
diff --git a/waku/persistence/driver_type.go b/waku/persistence/driver_type.go
new file mode 100644
index 00000000..86c784b4
--- /dev/null
+++ b/waku/persistence/driver_type.go
@@ -0,0 +1,22 @@
+package persistence
+
+import (
+	"database/sql"
+	"reflect"
+)
+
+const (
+	UndefinedDriver = iota
+	PostgresDriver
+	SQLiteDriver
+)
+
+func GetDriverType(db *sql.DB) int {
+	switch reflect.TypeOf(db.Driver()).String() {
+	case "*sqlite3.SQLiteDriver":
+		return SQLiteDriver
+	case "*stdlib.Driver":
+		return PostgresDriver
+	}
+	return UndefinedDriver
+}
diff --git a/waku/persistence/store.go b/waku/persistence/store.go
index 5964db8b..e9186856 100644
--- a/waku/persistence/store.go
+++ b/waku/persistence/store.go
@@ -5,7 +5,6 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
-	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -239,12 +238,13 @@ func (d *DBStore) cleanOlderRecords(ctx context.Context) error {
 
 	return nil
 }
+
 func (d *DBStore) getDeleteOldRowsQuery() string {
 	sqlStmt := `DELETE FROM message WHERE id IN (SELECT id FROM message ORDER BY receiverTimestamp DESC %s OFFSET $1)`
-	switch reflect.TypeOf(d.db.Driver()).String() {
-	case "*sqlite3.SQLiteDriver":
+	switch GetDriverType(d.db) {
+	case SQLiteDriver:
 		sqlStmt = fmt.Sprintf(sqlStmt, "LIMIT -1")
-	case "*stdlib.Driver":
+	case PostgresDriver:
 		sqlStmt = fmt.Sprintf(sqlStmt, "")
 	}
 	return sqlStmt
diff --git a/waku/persistence/utils/db.go b/waku/persistence/utils/db.go
index 96c5a919..ad7e3a95 100644
--- a/waku/persistence/utils/db.go
+++ b/waku/persistence/utils/db.go
@@ -6,6 +6,7 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/waku-org/go-waku/waku/persistence"
 	"github.com/waku-org/go-waku/waku/persistence/postgres"
 	"github.com/waku-org/go-waku/waku/persistence/sqlite"
 	"go.uber.org/zap"
@@ -65,3 +66,13 @@ func ExtractDBAndMigration(databaseURL string, dbSettings DBSettings, logger *za
 	return db, migrationFn, nil
 
 }
+
+func NewQueries(tbl string, db *sql.DB) (*persistence.Queries, error) {
+	switch persistence.GetDriverType(db) {
+	case persistence.SQLiteDriver:
+		return sqlite.NewQueries(tbl, db)
+	case persistence.PostgresDriver:
+		return postgres.NewQueries(tbl, db)
+	}
+	return nil, errors.New("unsupported database engine")
+}