Add more methods to zdb

Zondax · Oct 19, 2023 · dba4570 · dba4570
1 parent 17f8616
commit dba4570
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 3 deletions.
diff --git a/cmd/demo-metrics/main.go b/cmd/demo-metrics/main.go
@@ -2,14 +2,15 @@ package main
 
 import (
 	"github.com/zondax/golem/pkg/cli"
+	"github.com/zondax/golem/pkg/constants"
 	"github.com/zondax/golem/pkg/metrics"
 	"github.com/zondax/golem/pkg/runner"
 )
 
 func main() {
 	println("[Demo] Microservice example")
 
-	_, _ = cli.InitGlobalLogger(cli.DebugLevel)
+	_, _ = cli.InitGlobalLogger(constants.DebugLevel)
 
 	r := runner.NewRunner()
 

diff --git a/cmd/demo-panic/main.go b/cmd/demo-panic/main.go
@@ -2,14 +2,15 @@ package main
 
 import (
 	"github.com/zondax/golem/pkg/cli"
+	"github.com/zondax/golem/pkg/constants"
 	"github.com/zondax/golem/pkg/metrics"
 	"github.com/zondax/golem/pkg/runner"
 )
 
 func main() {
 	println("[Demo] Panic handler")
 
-	_, _ = cli.InitGlobalLogger(cli.DebugLevel)
+	_, _ = cli.InitGlobalLogger(constants.DebugLevel)
 
 	r := runner.NewRunner()
 

diff --git a/pkg/cli/root.go b/pkg/cli/root.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+	"github.com/zondax/golem/pkg/constants"
 	"go.uber.org/zap"
 	"os"
 )
@@ -69,7 +70,7 @@ func (c *CLI) init() {
 
 	// TODO: We can make this optional? and more configurable if we see the need
 	// Initialize logger
-	_, _ = InitGlobalLogger(DebugLevel)
+	_, _ = InitGlobalLogger(constants.DebugLevel)
 	setupCloseHandler(nil)
 	// Set Configuration Defaults
 	setupDefaultConfiguration(func() {

diff --git a/pkg/zdb/methods.go b/pkg/zdb/methods.go
@@ -58,6 +58,18 @@ func (z *zDatabase) Clauses(conds ...clause.Expression) ZDatabase {
 	return wrap(z.db.Clauses(conds...))
 }
 
+func (z *zDatabase) Select(query interface{}, args ...interface{}) ZDatabase {
+	return wrap(z.db.Select(query, args...))
+}
+
+func (z *zDatabase) Where(query interface{}, args ...interface{}) ZDatabase {
+	return wrap(z.db.Where(query, args...))
+}
+
+func (z *zDatabase) Limit(limit int) ZDatabase {
+	return wrap(z.db.Limit(limit))
+}
+
 func (z *zDatabase) Transaction(fc func(tx ZDatabase) error, opts ...*sql.TxOptions) (err error) {
 	return z.db.Transaction(func(tx *gorm.DB) error {
 		return fc(wrap(tx))

diff --git a/pkg/zdb/mock_zdb.go b/pkg/zdb/mock_zdb.go
@@ -52,6 +52,21 @@ func (m *MockZDatabase) Raw(sql string, values ...interface{}) ZDatabase {
 	return args.Get(0).(ZDatabase)
 }
 
+func (m *MockZDatabase) Select(query interface{}, values ...interface{}) ZDatabase {
+	args := m.Called(query, values[0])
+	return args.Get(0).(ZDatabase)
+}
+
+func (m *MockZDatabase) Where(query interface{}, values ...interface{}) ZDatabase {
+	args := m.Called(query, values[0])
+	return args.Get(0).(ZDatabase)
+}
+
+func (m *MockZDatabase) Limit(limit int) ZDatabase {
+	args := m.Called(limit)
+	return args.Get(0).(ZDatabase)
+}
+
 func (m *MockZDatabase) Exec(sql string, values ...interface{}) ZDatabase {
 	m.Called(sql, values)
 	return m

diff --git a/pkg/zdb/zdb.go b/pkg/zdb/zdb.go
@@ -22,6 +22,9 @@ type ZDatabase interface {
 	Scan(dest interface{}) ZDatabase
 	Rows() (*sql.Rows, error)
 	ScanRows(rows *sql.Rows, result interface{}) error
+	Select(query interface{}, args ...interface{}) ZDatabase
+	Where(query interface{}, args ...interface{}) ZDatabase
+	Limit(limit int) ZDatabase
 	Create(value interface{}) ZDatabase
 	Delete(value interface{}, where ...interface{}) ZDatabase
 	Raw(sql string, values ...interface{}) ZDatabase